Fix or ignore some pytype errors.

PiperOrigin-RevId: 675513581
jax-ml · Sep 17, 2024 · 7a703de · 7a703de
1 parent e59add0
commit 7a703de
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax_tpu_embedding/examples/singlehost_pjit_example.py b/jax_tpu_embedding/examples/singlehost_pjit_example.py
@@ -166,7 +166,7 @@ def build_embedding_configs(batch_size_per_device: int,
 
 
 def create_global_mesh(mesh_shape: Tuple[int, ...],
-                       axis_names: Sequence[jax.pxla.MeshAxisName]) -> Mesh:
+                       axis_names: Sequence[jax.pxla.MeshAxisName]) -> Mesh:  # pytype: disable=module-attr
   size = np.prod(mesh_shape)
   if len(jax.devices()) < size:
     raise ValueError(f'Test requires {size} global devices.')