code: def test_dynamic_shape_has_static_info_in_graph(
      self, i_shape, num_dims
  ):
      x_shape = (np.prod(i_shape[:num_dims]),) + (2, 2)
      y_shape = i_shape[:num_dims] + x_shape[1:]
  
      @tf.function
      def split(x, i, num_dims):
          y = batch_apply.split_leading_dim(x, i, num_dims)
          self.assertIsNotNone(y.shape.dims)
          y.shape.assert_is_compatible_with(y_shape)
          self.assertNotIn(None, y.shape[1:])
          return y
  
      # Build a concrete function with fully dynamic input dimensions.
      x = tf.TensorSpec((None,) + x_shape[1:])
      i = tf.TensorSpec((None,) + i_shape[1:])
      cf = split.get_concrete_function(x, i, num_dims)
      y_shape_dynamic = cf.output_shapes
      y_shape_dynamic.assert_is_compatible_with(y_shape)
  
imports: from absl.testing import parameterized
  import numpy
  from sonnet.src import base
  from sonnet.src import batch_apply
  from sonnet.src import test_utils
  import tensorflow
label_desc: Use image node_id 4 for calling the SplitLeadingDimTest obj's underlying member method code with example usage: obj.test_dynamic_shape_has_static_info_in_graph(i_shape, num_dims) and returns: y