code: def test_static_shape(self, i_shape, num_dims):
      x_shape = (np.prod(i_shape[:num_dims]),) + (2, 2)
      y_shape = i_shape[:num_dims] + x_shape[1:]
      x = tf.ones(x_shape)
      i = tf.ones(i_shape)
      y = batch_apply.split_leading_dim(x, i, num_dims)
      self.assertEqual(y.shape, 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 2 for calling the SplitLeadingDimTest obj's underlying member method code with example usage: obj.test_static_shape(i_shape, num_dims) without return types