Internal change.

PiperOrigin-RevId: 492381485

Author: fyangf

official/vision/modeling/layers/edgetpu.py

@@ -101,14 +101,15 @@ def _same(x):
   return tf.constant(1, dtype=x.dtype) + tf.math.floor(-x_clip)
 
 
-def shard_tensors(axis: int, block_size: int,
-                  *tensors: tf.Tensor) -> Iterable[Sequence[tf.Tensor]]:
+def shard_tensors(
+    axis: int, block_size: int,
+    tensors: Sequence[tf.Tensor]) -> Iterable[Sequence[tf.Tensor]]:
   """Consistently splits multiple tensors sharding-style.
 
   Args:
     axis: axis to be used to split tensors
     block_size: block size to split tensors.
-    *tensors: list of tensors.
+    tensors: list of tensors.
 
   Returns:
     List of shards, each shard has exactly one peace of each input tesnor.
@@ -211,7 +212,7 @@ def non_max_suppression_padded(boxes: tf.Tensor,
   scores = tf.reshape(scores, [batch_size, boxes_size])
   block = max(1, _RECOMMENDED_NMS_MEMORY // (boxes_size * boxes_size))
   indices = []
-  for boxes_i, scores_i in shard_tensors(0, block, boxes, scores):
+  for boxes_i, scores_i in shard_tensors(0, block, (boxes, scores)):
     indices.append(
         _non_max_suppression_as_is(boxes_i, scores_i, output_size,
                                    iou_threshold, refinements))

official/vision/modeling/layers/edgetpu_test.py

@@ -191,7 +191,7 @@ def test_shard_tensors(self):
         [15, 16, 17, 18, 19],
         [20, 21, 22, 23, 24],
     ]])
-    for i, (a_i, b_i) in enumerate(edgetpu.shard_tensors(1, 3, a, b)):
+    for i, (a_i, b_i) in enumerate(edgetpu.shard_tensors(1, 3, (a, b))):
       self.assertAllEqual(a_i, a[:, i * 3:i * 3 + 3])
       self.assertAllEqual(b_i, b[:, i * 3:i * 3 + 3, :])
 
@@ -233,7 +233,7 @@ def test_top_k_sharded_fusion_vs_top_k_unsharded(self, axis: int):
         shape=axis * [1] + [10000], dtype=tf.float32)
     top_1000_direct: tf.Tensor = tf.math.top_k(sample, 1000).values
     top_1000_sharded: Optional[tf.Tensor] = None
-    for (piece,) in edgetpu.shard_tensors(axis, 1500, sample):
+    for (piece,) in edgetpu.shard_tensors(axis, 1500, (sample,)):
       (top_1000_sharded,) = edgetpu.concat_and_top_k(
           1000, (top_1000_sharded, piece))
     self.assertAllEqual(top_1000_direct, top_1000_sharded)