Move _src/interpreters/ad.py to its own BUILD rule.

Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, and leads to improved build and iteration times.

This required moving some internal utilities out of dispatch.py, which is part of the main JAX build rule. I chose api_util.py because they seem to fit there.

PiperOrigin-RevId: 761722054

Author: Jake VanderPlas

jax/BUILD

@@ -316,7 +316,6 @@ py_library_providing_imports_info(
         "_src/ffi.py",
         "_src/flatten_util.py",
         "_src/interpreters/__init__.py",
-        "_src/interpreters/ad.py",
         "_src/interpreters/batching.py",
         "_src/interpreters/pxla.py",
         "_src/pjit.py",
@@ -381,6 +380,7 @@ py_library_providing_imports_info(
     visibility = ["//visibility:public"],
     deps = [
         ":abstract_arrays",
+        ":ad",
         ":ad_util",
         ":api_util",
         ":basearray",
@@ -671,6 +671,23 @@ pytype_strict_library(
     ] + py_deps("numpy"),
 )
 
+pytype_strict_library(
+    name = "ad",
+    srcs = ["_src/interpreters/ad.py"],
+    deps = [
+        ":ad_util",
+        ":api_util",
+        ":config",
+        ":core",
+        ":dtypes",
+        ":mesh",
+        ":partial_eval",
+        ":source_info_util",
+        ":tree_util",
+        ":util",
+    ],
+)
+
 pytype_strict_library(
     name = "mlir",
     srcs = ["_src/interpreters/mlir.py"],

jax/_src/api.py

@@ -37,6 +37,7 @@
 import numpy as np
 from contextlib import contextmanager
 
+from jax._src import api_util
 from jax._src import deprecations
 from jax._src import linear_util as lu
 from jax._src import stages
@@ -113,14 +114,14 @@ def _nan_check_posthook(fun, args, kwargs, output):
 
   try:
     dispatch.check_special(pjit.pjit_p.name, buffers)
-  except dispatch.InternalFloatingPointError as e:
+  except api_util.InternalFloatingPointError as e:
     assert config.debug_nans.value or config.debug_infs.value
     if hasattr(fun, '_fun'):
       f = fun._fun
       if getattr(f, '_apply_primitive', False):
         raise FloatingPointError(f"invalid value ({e.ty}) encountered in {f.__qualname__}") from None
       # compiled_fun can only raise in this case
-      dispatch.maybe_recursive_nan_check(e, f, args, kwargs)
+      api_util.maybe_recursive_nan_check(e, f, args, kwargs)
       raise AssertionError("Unreachable") from e
     else:
       # TODO(emilyaf): Shouldn't need this fallback.
@@ -1707,7 +1708,7 @@ def cache_miss(*args, **kwargs):
           out = execute(*p.flat_args)
         else:
           out = pxla.xla_pmap_p.bind_with_trace(trace, (p.flat_fun, *p.flat_args), params)
-      except dispatch.InternalFloatingPointError as e:
+      except api_util.InternalFloatingPointError as e:
         raise FloatingPointError(f'Invalid value ({e.ty}) encountered in parallel computation.')
 
     out_tree, out_flat = p.out_tree, out

jax/_src/api_util.py

@@ -767,3 +767,41 @@ def _check_no_aliased_closed_over_refs(dbg: core.DebugInfo, consts, args) -> Non
           f"array reference of type {a.str_short()} was both closed over and "
           f"passed as the argument "
           f"{dbg.safe_arg_names(len(args))[i]}" if dbg else "at flat index {i}")
+
+class InternalFloatingPointError(Exception):
+  name: str
+  ty: str
+
+  def __init__(self, name: str, ty: str):
+    self.name = name
+    self.ty = ty
+
+def maybe_recursive_nan_check(e: Exception, fun: Callable, args, kwargs,
+) -> None:  # always raises an exception
+  print("Invalid nan value encountered in the output of a jax.jit "
+        "function. Calling the de-optimized version.")
+  try:
+    _ = fun(*args, **kwargs)
+  except (FloatingPointError, ZeroDivisionError) as e2:
+    raise e2 from None
+  else:
+    _raise_no_nan_in_deoptimized(e)
+
+
+def _raise_no_nan_in_deoptimized(e) -> None:
+  msg = (f"{str(e)}. Because "
+        "jax_config.debug_nans.value and/or config.jax_debug_infs is set, the "
+        "de-optimized function (i.e., the function as if the `jit` "
+        "decorator were removed) was called in an attempt to get a more "
+        "precise error message. However, the de-optimized function did not "
+        "produce invalid values during its execution. This behavior can "
+        "result from `jit` optimizations causing the invalid value to be "
+        "produced. It may also arise from having nan/inf literals as "
+        "inputs or outputs, like `jax.jit(lambda ...: jax.numpy.nan)(...)`. "
+        "\n\n"
+        "It may be possible to avoid the invalid value by removing the "
+        "`jit` decorator, at the cost of losing optimizations. "
+        "\n\n"
+        "If you see this error, consider opening a bug report at "
+        "https://github.com/jax-ml/jax.")
+  raise FloatingPointError(msg) from None

jax/_src/dispatch.py

@@ -24,7 +24,7 @@
 import logging
 import threading
 import time
-from typing import Any, Callable
+from typing import Any
 
 import jax
 from jax._src import api
@@ -42,6 +42,7 @@
 from jax._src.interpreters import mlir
 from jax._src.interpreters import pxla
 from jax._src.interpreters import xla
+from jax._src.api_util import InternalFloatingPointError
 from jax._src.layout import DeviceLocalLayout, Layout
 from jax._src.lib import xla_client as xc
 from jax._src.mesh import AbstractMesh, Mesh
@@ -341,43 +342,6 @@ class CopySemantics(enum.Enum):
   COPY = enum.auto()
   DONATE = enum.auto()
 
-class InternalFloatingPointError(Exception):
-  name: str
-  ty: str
-
-  def __init__(self, name: str, ty: str):
-    self.name = name
-    self.ty = ty
-
-def maybe_recursive_nan_check(e: Exception, fun: Callable, args, kwargs,
-) -> None:  # always raises an exception
-  print("Invalid nan value encountered in the output of a jax.jit "
-        "function. Calling the de-optimized version.")
-  try:
-    _ = fun(*args, **kwargs)
-  except (FloatingPointError, ZeroDivisionError) as e2:
-    raise e2 from None
-  else:
-    _raise_no_nan_in_deoptimized(e)
-
-def _raise_no_nan_in_deoptimized(e) -> None:
-  msg = (f"{str(e)}. Because "
-        "jax_config.debug_nans.value and/or config.jax_debug_infs is set, the "
-        "de-optimized function (i.e., the function as if the `jit` "
-        "decorator were removed) was called in an attempt to get a more "
-        "precise error message. However, the de-optimized function did not "
-        "produce invalid values during its execution. This behavior can "
-        "result from `jit` optimizations causing the invalid value to be "
-        "produced. It may also arise from having nan/inf literals as "
-        "inputs or outputs, like `jax.jit(lambda ...: jax.numpy.nan)(...)`. "
-        "\n\n"
-        "It may be possible to avoid the invalid value by removing the "
-        "`jit` decorator, at the cost of losing optimizations. "
-        "\n\n"
-        "If you see this error, consider opening a bug report at "
-        "https://github.com/jax-ml/jax.")
-  raise FloatingPointError(msg) from None
-
 def _identity_fn(x):
   return x
 

jax/_src/interpreters/ad.py

@@ -21,12 +21,12 @@
 from functools import partial
 from typing import Any
 
+from jax._src import api_util
 from jax._src import config
-from jax._src import dispatch
 from jax._src import linear_util as lu
 from jax._src.interpreters import partial_eval as pe
-from jax.tree_util import (tree_flatten, tree_unflatten,
-                           register_pytree_node, Partial, PyTreeDef)
+from jax._src.tree_util import (tree_flatten, tree_unflatten,
+                                register_pytree_node, Partial, PyTreeDef)
 from jax._src import mesh as mesh_lib
 from jax._src import core
 from jax._src import source_info_util
@@ -1125,7 +1125,7 @@ def out_axes_thunk():
 
   try:
     out_flat = primitive.bind(fun, *all_args, **new_params)
-  except dispatch.InternalFloatingPointError as e:
+  except api_util.InternalFloatingPointError as e:
     print("Invalid nan value encountered in the backward pass of a jax.jit "
           "function. Calling the de-optimized backward pass.")
     try:
@@ -1135,7 +1135,7 @@ def out_axes_thunk():
     else:
       # If control reaches this line, we got a NaN on the output of `compiled`
       # but not `fun.call_wrapped` on the same arguments. Let's tell the user.
-      dispatch._raise_no_nan_in_deoptimized(e)
+      api_util._raise_no_nan_in_deoptimized(e)
   arg_cts = tree_unflatten(out_tree(), out_flat)
 
   # The freevars are being fanned out (not mapped). During transpose the
@@ -1266,11 +1266,3 @@ def __init__(self):
 
 # TODO(mattjj): remove this vestigial dict
 reducing_transposes: dict[core.Primitive, Callable] = {}
-
-########################### pvary ##################################
-
-def _pvary_transpose_rule(cts, *_, axes, axis_index_groups):
-  from jax._src.lax import parallel as lax_parallel
-  return lax_parallel.psum_invariant_p.bind(
-      *cts, axes=axes, axis_index_groups=axis_index_groups)
-deflinear2(core.pvary_p, _pvary_transpose_rule)

jax/_src/lax/parallel.py

@@ -2059,3 +2059,10 @@ def _psum_invariant_transpose_rule(cts, *args, axes, axis_index_groups):
   del args
   return core.pvary_p.bind(*cts, axes=axes, axis_index_groups=axis_index_groups)
 ad.deflinear2(psum_invariant_p, _psum_invariant_transpose_rule)
+
+########################### pvary ##################################
+
+def _pvary_transpose_rule(cts, *_, axes, axis_index_groups):
+  return psum_invariant_p.bind(
+      *cts, axes=axes, axis_index_groups=axis_index_groups)
+ad.deflinear2(core.pvary_p, _pvary_transpose_rule)

jax/_src/pjit.py

@@ -176,10 +176,10 @@ def _python_pjit_helper(fun: Callable, jit_info: PjitInfo, *args, **kwargs):
               f"Argument '{name}' of shape {aval.str_short()} of type"
               f' {type(arg)} is not a valid JAX type.') from e
       raise AssertionError("Unreachable") from e
-  except dispatch.InternalFloatingPointError as e:
+  except api_util.InternalFloatingPointError as e:
     if getattr(fun, '_apply_primitive', False):
       raise FloatingPointError(f"invalid value ({e.ty}) encountered in {fun.__qualname__}") from None
-    dispatch.maybe_recursive_nan_check(e, fun, args, kwargs)
+    api_util.maybe_recursive_nan_check(e, fun, args, kwargs)
 
   if p.box_data:
     box_treedef, out_tree = p.out_tree.children()
@@ -2562,7 +2562,7 @@ def prune_type(ty, xs, maybe_zeros):
         keep_unused=keep_unused,
         inline=inline,
         compiler_options_kvs=compiler_options_kvs)
-  except dispatch.InternalFloatingPointError as e:
+  except api_util.InternalFloatingPointError as e:
     print("Invalid nan value encountered in the backward pass of a jax.jit "
           "function. Calling the de-optimized backward pass.")
     try:
@@ -2572,7 +2572,7 @@ def prune_type(ty, xs, maybe_zeros):
     else:
       # If control reaches this line, we got a NaN on the output of `compiled`
       # but not `fun.call_wrapped` on the same arguments. Let's tell the user.
-      dispatch._raise_no_nan_in_deoptimized(e)
+      api_util._raise_no_nan_in_deoptimized(e)
 
   if attrs_tracked:
     final_states, nz_cts_out = split_list(nz_cts_out, [num_attr_outs])

jax/_src/shard_map.py

@@ -1054,7 +1054,7 @@ def _maybe_check_special(outs):
           for s in getattr(leaf, 'addressable_shards', [])]
   try:
     dispatch.check_special('shard_map', bufs)
-  except dispatch.InternalFloatingPointError as e:
+  except api_util.InternalFloatingPointError as e:
     raise FloatingPointError(f'Invalid value ({e.ty}) encountered in sharded computation.') from None
 
 class ShardMapTrace(core.Trace):
@@ -1562,7 +1562,7 @@ def new_out_specs_thunk():
     except (FloatingPointError, ZeroDivisionError) as e2:
       raise e2 from None
     else:
-      dispatch._raise_no_nan_in_deoptimized(e)
+      api_util._raise_no_nan_in_deoptimized(e)
   return tree_unflatten(out_tree(), out_flat)
 ad.primitive_transposes[shard_map_p] = _shard_map_transpose
 

jax/extend/BUILD

@@ -43,6 +43,7 @@ py_library_providing_imports_info(
     deps = [
         "//jax",
         "//jax:abstract_arrays",
+        "//jax:ad",
         "//jax:ad_util",
         "//jax:core",
     ],