code: def test_unshuffle_stacking_basic(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)
      stacked_shards = [
          array_ops.concat([i, j], axis=0)
          for i, j in zip(x_shards, y_shards)
      ]
      stacked = array_ops.concat(stacked_shards, axis=0)
      self.assertAllEqual(
          v3_utils.unshuffle_from_sc_to_cpu(
              t=stacked,
              num_sparse_cores=num_sc,
              offset_in_shard=0,
              size_in_shard=ta.vocab // num_sc,
              shard_rotation=ta.shift,
          ),
          x,
      )
      self.assertAllEqual(
          v3_utils.unshuffle_from_sc_to_cpu(
              t=stacked,
              num_sparse_cores=num_sc,
              offset_in_shard=ta.vocab // num_sc,
              size_in_shard=tb.vocab // num_sc,
              shard_rotation=tb.shift,
          ),
          y,
      )
  
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 3 for calling the TpuEmbeddingV3UtilsTest obj's underlying member method code with example usage: obj.test_unshuffle_stacking_basic() without return types