code: def test_shapes(
      batch_shape: Tuple,
      num_pieces: int,
      num_samples: int,
      serialize_fn,
  ):
      gamma = mx.nd.ones(shape=(*batch_shape,))
      slopes = mx.nd.ones(
          shape=(*batch_shape, num_pieces)
      )  # all positive
      knot_spacings = (
          mx.nd.ones(shape=(*batch_shape, num_pieces)) / num_pieces
      )  # positive and sum to 1
      target = mx.nd.ones(shape=batch_shape)  # shape of gamma
  
      distr = PiecewiseLinear(
          gamma=gamma, slopes=slopes, knot_spacings=knot_spacings
      )
      distr = serialize_fn(distr)
  
      # 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, axis=-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, axis=0).shape == (
          num_samples,
          *batch_shape,
      )
  
file_path: gluonts/test/mx/distribution/test_piecewise_linear.py
imports: from typing import List, Tuple
  import mxnet
  import numpy
  import pytest
  from gluonts.core.serde import dump_json, load_json
  from gluonts.mx.distribution import FixedKnotsPiecewiseLinearOutput, PiecewiseLinear, PiecewiseLinearOutput
  from gluonts.testutil import empirical_cdf
label_desc: Use image node_id 2 for calling a global function with example usage: test_shapes(batch_shape, num_pieces, num_samples, serialize_fn) without return types