code: def testSimpleTwoDimensional(self):
      with self.session():
          data = constant_op.constant(
              [
                  [0, 1, 2],
                  [3, 4, 5],
                  [6, 7, 8],
                  [9, 10, 11],
                  [12, 13, 14],
                  [15, 16, 17],
              ],
              dtype=dtypes.float32,
          )
          indices = constant_op.constant([0, 0, 2, 3, 2, 1])
          partitions = data_flow_ops.dynamic_partition(
              data, indices, num_partitions=4
          )
          partition_vals = self.evaluate(partitions)
  
      self.assertEqual(4, len(partition_vals))
      self.assertAllEqual([[0, 1, 2], [3, 4, 5]], partition_vals[0])
      self.assertAllEqual([[15, 16, 17]], partition_vals[1])
      self.assertAllEqual([[6, 7, 8], [12, 13, 14]], partition_vals[2])
      self.assertAllEqual([[9, 10, 11]], partition_vals[3])
      # Vector data input to DynamicPartition results in
      # `num_partitions` matrices with an unknown number of rows, and 3 columns.
      self.assertEqual([None, 3], partitions[0].get_shape().as_list())
      self.assertEqual([None, 3], partitions[1].get_shape().as_list())
      self.assertEqual([None, 3], partitions[2].get_shape().as_list())
      self.assertEqual([None, 3], partitions[3].get_shape().as_list())
  
imports: import unittest
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import data_flow_ops
  from tensorflow.python.ops import gradients_impl
  import tensorflow.python.ops.data_flow_grad
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the DynamicPartitionTest obj's underlying member method code with example usage: obj.testSimpleTwoDimensional() without return types