code: def test_shapes(
      batch_shape: Tuple, num_pieces: int, num_samples: int
  ):
      gamma = torch.ones(size=(*batch_shape,))
      slopes = torch.ones(
          size=(*batch_shape, num_pieces)
      )  # all positive
      knot_spacings = (
          torch.ones(size=(*batch_shape, num_pieces)) / num_pieces
      )  # positive and sum to 1
      target = torch.ones(size=batch_shape)  # shape of gamma
  
      distr = PiecewiseLinear(
          gamma=gamma, slopes=slopes, knot_spacings=knot_spacings
      )
  
      # assert that the parameters and target have proper shapes
      assert gamma.shape == target.shape
      assert knot_spacings.shape == slopes.shape
      assert len(gamma.shape) + 1 == len(knot_spacings.shape)
  
      # assert that batch_shape is computed properly
      assert distr.batch_shape == batch_shape
  
      # assert that shapes of original parameters are correct
      assert distr.b.shape == slopes.shape
      assert distr.knot_positions.shape == knot_spacings.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 a_tilde
      assert distr.quantile_internal(knot_spacings, dim=-2).shape == (
          *batch_shape,
          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_piecewise_linear.py
imports: from typing import Tuple, List
  import pytest
  import torch
  import numpy
  from gluonts.torch.distributions import PiecewiseLinear, PiecewiseLinearOutput
label_desc: Use image node_id 3 for calling a global function with example usage: test_shapes(batch_shape, num_pieces, num_samples) without return types