code: def testDifferentShapesBetweenCalls(self):
      self.skipForTfrt(
          "b/269333905, TFRT cpu fails due to step_id not propagated."
      )
      self.skipForDeviceType(
          ["TPU"],
          "Known failure under TPU for legalization requires a static shape.",
      )
  
      # The error only happens across the batch, where the value of
      # tf.unique are differnet.
      def produce_data(inputs, label):
          inputs = api.relayout(
              inputs, Layout.batch_sharded(self.mesh, _MESH_DIM_X, 1)
          )
          label = api.relayout(
              label, Layout.batch_sharded(self.mesh, _MESH_DIM_X, 1)
          )
          return inputs, label
  
      @polymorphic_function.function
      def train_fn(inputs, label):
          inputs, indices = array_ops.unique(inputs)
          return math_ops.unsorted_segment_sum(
              label, indices, len(inputs)
          )
  
      input1, label1 = produce_data([6, 0, 6, 0], [1, 2, 3, 4])
      input2, label2 = produce_data([2, 1, 2, 0], [1, 2, 3, 4])
  
      result1 = train_fn(input1, label1)
      result2 = train_fn(input2, label2)
      self.assertAllEqual(
          result1.numpy(),
          [
              4,
              6,
          ],
      )
      self.assertAllEqual(
          result2.numpy(),
          [
              4,
              2,
              4,
          ],
      )
  
imports: import os
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import dtensor_device
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 13 for calling the CollectiveTest obj's underlying member method code with example usage: obj.testDifferentShapesBetweenCalls() and returns: inputs,  label, math_ops