code: def testValues(self):
      segment_shapes = [tf.TensorShape([2]), tf.TensorShape([3])]
      inputs = tf.constant(
          np.hstack([np.zeros((5, 2)), np.ones((5, 3))]),
          dtype=tf.float32,
      )
      first, second = util.segment_dim(
          inputs, dim=1, shapes=segment_shapes
      )
  
      self.assertAllEqual(first.numpy(), np.zeros_like(first))
      self.assertAllEqual(second.numpy(), np.ones_like(second))
  
imports: from absl.testing import parameterized
  import numpy
  from sonnet.src import linear
  from sonnet.src import test_utils
  from sonnet.src.nets.dnc import util
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling the SegmentDimTest obj's underlying member method code with example usage: obj.testValues() without return types