code: def test_index_mapping_stacked_tables(self):
      num_sc = 4
      ta = TestTable(vocab=12, dim=4, shift=0)
      tb = TestTable(vocab=32, dim=4, shift=1)
      x, x_shards = create_test_table_shards(ta, num_sc)
      y, y_shards = create_test_table_shards(
          tb, num_sc, table_data_start=100
      )
      stacked_shards = [
          array_ops.concat([i, j], axis=0)
          for i, j in zip(x_shards, y_shards)
      ]
      indices_ta = tf_constant([1, 2, 7, 9, 11])
      shard_idx, position_in_shard = v3_utils.map_indices_in_shard(
          num_sparse_cores=num_sc,
          offset_in_shard=0,
          shard_rotation=ta.shift,
          row_indices=indices_ta,
      )
      self.assertAllEqual(
          [x[i] for i in indices_ta],
          [
              stacked_shards[j][k]
              for j, k in zip(shard_idx, position_in_shard)
          ],
      )
      indices_tb = tf_constant([1, 2, 7, 9, 15, 27])
      shard_idx, position_in_shard = v3_utils.map_indices_in_shard(
          num_sparse_cores=num_sc,
          offset_in_shard=ta.vocab // num_sc,
          shard_rotation=tb.shift,
          row_indices=indices_tb,
      )
      self.assertAllEqual(
          [y[i] for i in indices_tb],
          [
              stacked_shards[j][k]
              for j, k in zip(shard_idx, position_in_shard)
          ],
      )
  
file_path: tensorflow-upstream/tensorflow/python/tpu/tpu_embedding_v3_utils_test.py
imports: import collections
  from absl.testing import parameterized
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.framework.constant_op import constant
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import tpu_embedding_v3_utils
label_desc: Use image node_id 6 for calling the TpuEmbeddingV3UtilsTest obj's underlying member method code with example usage: obj.test_index_mapping_stacked_tables() without return types