code: def test_shapes(
      batch_shape: Tuple,
      num_pieces: int,
      num_qk: int,
      num_samples: int,
  ):
      spline_knots = torch.ones(
          (*batch_shape, (num_qk - 1), num_pieces), dtype=torch.float32
      )
      spline_heights = torch.ones(
          (*batch_shape, (num_qk - 1), num_pieces), dtype=torch.float32
      )
      beta_l = torch.ones(batch_shape, dtype=torch.float32)
      beta_r = torch.ones(batch_shape, dtype=torch.float32)
  
      # quantile knot positions are non-decreasing
      qk_y = torch.cumsum(
          torch.ones((*batch_shape, num_qk), dtype=torch.float32),
          dim=-1,
      )
      qk_x = torch.cumsum(
          torch.ones((*batch_shape, num_qk), dtype=torch.float32),
          dim=-1,
      )
  
      target = torch.ones(batch_shape, dtype=torch.float32)
  
      distr = ISQF(
          spline_knots=spline_knots,
          spline_heights=spline_heights,
          beta_l=beta_l,
          beta_r=beta_r,
          qk_y=qk_y,
          qk_x=qk_x,
      )
  
      # assert that the parameters and target have proper shapes
      assert spline_knots.shape == spline_heights.shape
      assert beta_l.shape == beta_r.shape == target.shape
      assert qk_y.shape == qk_x.shape
      assert len(qk_y.shape) == len(target.shape) + 1
      assert len(spline_knots.shape) == len(target.shape) + 2
  
      # assert that batch_shape is computed properly
      assert distr.batch_shape == batch_shape
  
      # assert that shapes of original parameters are correct
      assert distr.sk_x.shape == distr.sk_y.shape == spline_knots.shape
      assert (
          distr.delta_sk_x.shape
          == distr.delta_sk_y.shape
          == spline_knots.shape
      )
  
      assert (
          distr.qk_x.shape
          == distr.qk_y.shape
          == spline_knots.shape[:-1]
      )
      assert (
          distr.qk_x_plus.shape
          == distr.qk_y_plus.shape
          == spline_knots.shape[:-1]
      )
      assert distr.qk_x_l.shape == distr.qk_y_l.shape == batch_shape
      assert distr.qk_x_r.shape == distr.qk_y_r.shape == batch_shape
  
      assert distr.tail_al.shape == distr.tail_bl.shape == batch_shape
      assert distr.tail_ar.shape == distr.tail_br.shape == batch_shape
  
      # assert that the shape of crps is correct
      assert distr.crps(target).shape == batch_shape
  
      # assert that the quantile shape is correct when computing
      # the quantile values at knot positions - used for alpha_tilde
      assert distr.quantile_spline(spline_knots, dim=-2).shape == (
          *batch_shape,
          num_qk - 1,
          num_pieces,
      )
  
      # assert that the samples and the quantile values shape when num_samples is None is correct
      samples = distr.sample()
      assert samples.shape == batch_shape
      assert distr.quantile_internal(samples).shape == batch_shape
  
      # assert that the samples and the quantile values shape when num_samples is not None is correct
      samples = distr.sample((num_samples,))
      assert samples.shape == (num_samples, *batch_shape)
      assert distr.quantile_internal(samples, dim=0).shape == (
          num_samples,
          *batch_shape,
      )
  
file_path: gluonts/test/torch/distribution/test_torch_isqf.py
imports: from typing import List, Tuple
  import torch
  import numpy
  import pytest
  from gluonts.torch.distributions import ISQF, ISQFOutput
label_desc: Use image node_id 2 for calling a global function with example usage: test_shapes(batch_shape, num_pieces, num_qk, num_samples) without return types