code: def testRaggedSegmentStack(
      self,
      data,
      partitions,
      num_partitions,
      expected,
      data_ragged_rank=None,
      segment_ids_ragged_rank=None,
      expected_ragged_rank=None,
  ):
      for seg_dtype in [dtypes.int32, dtypes.int64]:
          data_tensor = ragged_factory_ops.constant(
              data,
              row_splits_dtype=seg_dtype,
              ragged_rank=data_ragged_rank,
          )
          segment_ids_tensor = ragged_factory_ops.constant(
              partitions,
              dtype=seg_dtype,
              row_splits_dtype=seg_dtype,
              ragged_rank=segment_ids_ragged_rank,
          )
          expected_tensor = ragged_factory_ops.constant(
              expected,
              row_splits_dtype=seg_dtype,
              ragged_rank=expected_ragged_rank,
          )
          result = ragged_array_ops.stack_dynamic_partitions(
              data_tensor, segment_ids_tensor, num_partitions
          )
          self.assertAllEqual(result, expected_tensor)
  
          # Check that it's equivalent to tf.stack(dynamic_partition(...)),
          # where applicable.
          if (
              data_ragged_rank == 0
              and segment_ids_ragged_rank == 0
              and seg_dtype == dtypes.int32
          ):
              equiv = ragged_concat_ops.stack(
                  data_flow_ops.dynamic_partition(
                      data_tensor, segment_ids_tensor, num_partitions
                  )
              )
              self.assertAllEqual(
                  result, self.evaluate(equiv).to_list()
              )
  
file_path: tensorflow-upstream/tensorflow/python/ops/ragged/ragged_dynamic_partition_op_test.py
imports: from absl.testing import parameterized
  from tensorflow.python.eager import context
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  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.ragged import ragged_array_ops
  from tensorflow.python.ops.ragged import ragged_concat_ops
  from tensorflow.python.ops.ragged import ragged_factory_ops
  from tensorflow.python.platform import googletest
label_desc: Use image node_id 1 for calling the RaggedSegmentStackOpTest obj's underlying member method code with example usage: obj.testRaggedSegmentStack(data, partitions, num_partitions, expected, data_ragged_rank, segment_ids_ragged_rank, expected_ragged_rank) without return types