Automated Code Change

PiperOrigin-RevId: 753039040

Author: h-joo

jax/_src/core.py

@@ -2047,7 +2047,7 @@ def standard_insert_pvary(*args):
   if not args:
     return args
   in_vma = [frozenset() if (aval := get_aval(a)) is abstract_token
-            else aval.vma for a in args]
+            else aval.vma for a in args]  # pytype: disable=attribute-error
   out_vma = frozenset.union(*in_vma)
   return [pvary(arg, tuple(n for n in out_vma if n not in src))
           if out_vma - src else arg for arg, src in zip(args, in_vma)]

jax/_src/interpreters/mlir.py

@@ -1097,7 +1097,7 @@ class UnconstrainedVariants(NamedTuple):
 
 def _get_unconstrained_variants(s, aval) -> UnconstrainedVariants:
   us = contains_unconstrained(s)
-  unconstrained_dims = ({i for i, p in enumerate(s.spec)
+  unconstrained_dims = ({i for i, p in enumerate(s.spec)  # pytype: disable=attribute-error
                          if p is PartitionSpec.UNCONSTRAINED} if us else None)
   return UnconstrainedVariants(
       contains_unconstrained=us, all_unconstrained=all_unconstrained(s, aval),

jax/_src/lax/lax.py

@@ -5172,7 +5172,7 @@ def _dot_general_shape_computation(lhs_shape, rhs_shape, dimension_numbers):
   lhs_tensored_shape = tuple_delete(lhs_shape, lhs_contract_or_batch)
   rhs_group = ()
   if isinstance(dimension_numbers, RaggedDotDimensionNumbers):
-    rhs_group = tuple(dimension_numbers.rhs_group_dimensions)
+    rhs_group = tuple(dimension_numbers.rhs_group_dimensions)  # pytype: disable=attribute-error
   rhs_contract_or_batch_or_group = tuple(
       sorted(tuple(rhs_contracting) + tuple(rhs_batch) + rhs_group)
   )
@@ -6017,7 +6017,7 @@ def grad_x_dims():
         unsorted_axes = list(x_batch) + x_kept + x_contract_sorted_by_y
       case RaggedDotMode.RAGGED_CONTRACTING | RaggedDotMode.RAGGED_BATCH:
         raise unimplemented('grad_x_dims', mode)
-    return dims, unsorted_axes
+    return dims, unsorted_axes  # pytype: disable=name-error
 
   def grad_y_dims():
     match mode:
@@ -6036,7 +6036,7 @@ def grad_y_dims():
         )
       case RaggedDotMode.RAGGED_CONTRACTING | RaggedDotMode.RAGGED_BATCH:
         raise unimplemented('grad_y_dims', mode)
-    return dims, unsorted_axes
+    return dims, unsorted_axes  # pytype: disable=name-error
 
   def _ragged_dot_grad(lhs, rhs, dims_fn, aval):
     dims, unsorted_axes = dims_fn()
@@ -6238,7 +6238,7 @@ def expand(x, dim, gs, *axes):
           lhs,
           rhs,
           dimension_numbers=ragged_dot_dimension_numbers.dot_dimension_numbers,
-      )
+      )  # pytype: disable=bad-return-type
 
 
 def _ragged_dot_general_lower(

jax/_src/pallas/triton/pallas_call_registration.py

@@ -40,7 +40,7 @@ def normalize_grid(grid: pallas_core.StaticGrid) -> tuple[int, int, int]:
 
 
 def avals_to_layouts(avals):
-  return [list(reversed(range(aval.ndim))) for aval in avals]
+  return [list(reversed(range(aval.ndim))) for aval in avals]  # pytype: disable=attribute-error
 
 
 def pallas_call_lowering(

jax/_src/pjit.py

@@ -750,16 +750,16 @@ def _infer_input_type(fun: Callable, dbg: core.DebugInfo,
     for i, x in enumerate(explicit_args):
       avals.append(core.shaped_abstractify(x))
   except OverflowError:
-    arg_path = f"argument path is {dbg.arg_names[i]}"
+    arg_path = f"argument path is {dbg.arg_names[i]}"  # pytype: disable=name-error
     raise OverflowError(
       "An overflow was encountered while parsing an argument to a jitted "
       f"computation, whose {arg_path}."
     ) from None
   except TypeError:
-    arg_description = f"path {dbg.arg_names[i]}"
+    arg_description = f"path {dbg.arg_names[i]}"  # pytype: disable=name-error
     raise TypeError(
       f"Error interpreting argument to {fun} as an abstract array."
-      f" The problematic value is of type {type(x)} and was passed to"
+      f" The problematic value is of type {type(x)} and was passed to"  # pytype: disable=name-error
       f" the function at {arg_description}.\n"
       "This typically means that a jit-wrapped function was called with a non-array"
       " argument, and this argument was not marked as static using the"
@@ -2035,8 +2035,8 @@ def _pjit_cached_lower_jaxpr_to_fun(ctx: mlir.LoweringRuleContext,
   elif isinstance(axis_ctx, sharding_impls.SPMDAxisContext):
     num_devices = axis_ctx.mesh.size
   key = (pjit_p, name, jaxpr, effects, num_devices,
-         pxla.SemanticallyEqualShardings(in_shardings, jaxpr.in_avals),
-         pxla.SemanticallyEqualShardings(out_shardings, jaxpr.out_avals),
+         pxla.SemanticallyEqualShardings(in_shardings, jaxpr.in_avals),  # pytype: disable=wrong-arg-types
+         pxla.SemanticallyEqualShardings(out_shardings, jaxpr.out_avals),  # pytype: disable=wrong-arg-types
          in_layouts, out_layouts, api_name)
 
   func = mod_ctx.cached_primitive_lowerings.get(key, None)

jax/_src/tpu_custom_call.py

@@ -222,7 +222,7 @@ def _tpu_custom_call_abstract_eval(*_, out_avals, **__):
 
 
 def _avals_to_layouts(avals) -> Sequence[Sequence[int]]:
-  return [tuple(range(a.ndim - 1, -1, -1)) for a in avals]
+  return [tuple(range(a.ndim - 1, -1, -1)) for a in avals]  # pytype: disable=attribute-error
 
 
 def _tpu_custom_call_lowering(