code: def _safe_split(x, num_outputs, axis, squeeze_axis, *args, **kwargs):
      """
      A type-stable wrapper around mx.nd.split.
  
      Currently mx.nd.split behaves weirdly if num_outputs=1:
  
          a = mx.nd.ones(shape=(1, 1, 2))
          l = a.split(axis=1, num_outputs=1, squeeze_axis=False)
          type(l)  # mx.NDArray
          l.shape  # (1, 1, 2)
  
      Compare that with the case num_outputs=2:
  
          a = mx.nd.ones(shape=(1, 2, 2))
          l = a.split(axis=1, num_outputs=2, squeeze_axis=False)
          type(l)  # list
          len(l)  # 2
          l[0].shape  # (1, 1, 2)
  
      This wrapper makes the behavior consistent by always returning a list
      of length num_outputs, whose elements will have one less axis than x
      in case x.shape[axis]==num_outputs and squeeze_axis==True, and the same
      number of axes as x otherwise.
      """
      if num_outputs > 1:
          return x.split(
              axis=axis,
              num_outputs=num_outputs,
              squeeze_axis=squeeze_axis,
              *args,
              **kwargs,
          )
      return [x.squeeze(axis=axis)] if squeeze_axis else [x]
  
imports: from typing import Optional, Tuple
  import mxnet
  import numpy
  from gluonts.core.component import validated
  from gluonts.mx import Tensor
  from gluonts.mx.linalg_util import jitter_cholesky
  from gluonts.mx.util import _broadcast_param, make_nd_diag
  from .None import Distribution, Gaussian, MultivariateGaussian
  from .distribution import getF
label_desc: Use image node_id 1 for calling a global function with example usage: _safe_split(x, num_outputs, axis, squeeze_axis) and returns: unknown, x