code: def test_jitted_slices(self):
      self.skip_if_no_xla()
      shp = [7, 6, 5, 4]
      t = tf.cast(tf.reshape(tf.range(np.prod(shp)), shp), tf.float32)
  
      @tf.function(jit_compile=True)
      def f(ix):
          return slicing._slice_params_to_dict(
              mvn_diag.MultivariateNormalDiag(t, tf.ones([shp[-1]])),
              slices=make_slices[..., ix, :],
          )
  
      self.assertAllEqual(t[:, 3], f(tf.constant(3))["loc"])
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import scale_matvec_diag
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import split
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import slicing
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 20 for calling the SlicingTest obj's underlying member method code with example usage: obj.test_jitted_slices() and returns: slicing