Internal fix.

PiperOrigin-RevId: 493182471

Author: fyangf

official/vision/modeling/layers/edgetpu.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 
 """EdgeTPU oriented layers and tools."""
-from collections.abc import Iterable, Sequence
-from typing import List, Optional, Union
+from typing import Optional, Sequence, Union, Iterable
 
 import numpy as np
 import tensorflow as tf
@@ -102,8 +101,8 @@ def _same(x):
 
 
 def shard_tensors(
-    axis: int, block_size: int,
-    tensors: Sequence[tf.Tensor]) -> Iterable[Sequence[tf.Tensor]]:
+    axis: int, block_size: int, tensors: Sequence[tf.Tensor]
+) -> Union[list[Sequence[tf.Tensor]], Iterable[Sequence[tf.Tensor]]]:
   """Consistently splits multiple tensors sharding-style.
 
   Args: