Merge pull request #29186 from jakevdp:jax-numpy-imports

PiperOrigin-RevId: 766730509

Author: Google-ML-Automation

jax/_src/lax/fft.py

@@ -21,15 +21,14 @@
 
 import numpy as np
 
-from jax import lax
-
 from jax._src import dispatch
 from jax._src import dtypes
 from jax._src.api import jit, linear_transpose, ShapeDtypeStruct
 from jax._src.core import Primitive, is_constant_shape
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
+from jax._src.lax import lax
 from jax._src.lib.mlir.dialects import hlo
 
 __all__ = [

jax/_src/numpy/array_api_metadata.py

@@ -21,7 +21,6 @@
 
 from types import ModuleType
 
-import jax
 from jax._src.sharding import Sharding
 from jax._src.lib import xla_client as xc
 from jax._src import config
@@ -40,6 +39,7 @@ def __array_namespace__(self, *, api_version: None | str = None) -> ModuleType:
   if api_version is not None and api_version != __array_api_version__:
     raise ValueError(f"{api_version=!r} is not available; "
                      f"available versions are: {[__array_api_version__]}")
+  import jax.numpy  # pytype: disable=import-error
   return jax.numpy
 
 
@@ -77,7 +77,7 @@ def default_device(self):
   def devices(self):
     out = [None]  # None indicates "uncommitted"
     for backend in xb.backends():
-        out.extend(jax.devices(backend))
+        out.extend(xb.devices(backend))
     return out
 
   def capabilities(self):

jax/_src/numpy/array_creation.py

@@ -19,18 +19,16 @@
 
 import numpy as np
 
-import jax
-from jax import lax
-from jax._src.api import jit
+from jax._src.api import device_put, jit
 from jax._src import core
 from jax._src import dtypes
-from jax._src.lax import lax as lax_internal
+from jax._src.lax import lax
 from jax._src.lib import xla_client as xc
 from jax._src.numpy import ufuncs
 from jax._src.numpy import util
+from jax._src.sharding import Sharding
 from jax._src.typing import Array, ArrayLike, DuckTypedArray, DTypeLike
 from jax._src.util import canonicalize_axis, set_module
-from jax.sharding import Sharding
 
 
 export = set_module('jax.numpy')
@@ -205,15 +203,17 @@ def full(shape: Any, fill_value: ArrayLike,
     Array([[0, 1, 2],
            [0, 1, 2]], dtype=int32)
   """
+  from jax._src.numpy.lax_numpy import asarray  # pytype: disable=import-error
+
   dtypes.check_user_dtype_supported(dtype, "full")
   util.check_arraylike("full", fill_value)
 
   if np.ndim(fill_value) == 0:
     shape = canonicalize_shape(shape)
     return lax.full(shape, fill_value, dtype, sharding=util.normalize_device_to_sharding(device))
   else:
-    return jax.device_put(
-        util._broadcast_to(jax.numpy.asarray(fill_value, dtype=dtype), shape), device)
+    return device_put(
+        util._broadcast_to(asarray(fill_value, dtype=dtype), shape), device)
 
 
 @export
@@ -394,6 +394,8 @@ def full_like(a: ArrayLike | DuckTypedArray,
     Array([[1, 1, 1],
            [2, 2, 2]], dtype=int32)
   """
+  from jax._src.numpy.lax_numpy import asarray  # pytype: disable=import-error
+
   if hasattr(a, 'dtype') and hasattr(a, 'shape'):  # support duck typing
     util.check_arraylike("full_like", 0, fill_value)
   else:
@@ -408,8 +410,8 @@ def full_like(a: ArrayLike | DuckTypedArray,
   else:
     shape = np.shape(a) if shape is None else shape  # type: ignore[arg-type]
     dtype = dtypes.result_type(a) if dtype is None else dtype
-    return jax.device_put(
-        util._broadcast_to(jax.numpy.asarray(fill_value, dtype=dtype), shape), device)
+    return device_put(
+        util._broadcast_to(asarray(fill_value, dtype=dtype), shape), device)
 
 @overload
 def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
@@ -510,6 +512,8 @@ def _linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
               axis: int = 0,
               *, device: xc.Device | Sharding | None = None) -> Array | tuple[Array, Array]:
   """Implementation of linspace differentiable in start and stop args."""
+  from jax._src.numpy.lax_numpy import asarray  # pytype: disable=import-error
+
   dtypes.check_user_dtype_supported(dtype, "linspace")
   if num < 0:
     raise ValueError(f"Number of samples, {num}, must be non-negative.")
@@ -529,13 +533,13 @@ def _linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
   bounds_shape.insert(axis, 1)
   div = (num - 1) if endpoint else num
   if num > 1:
-    delta: Array = lax.convert_element_type(stop - start, computation_dtype) / jax.numpy.array(div, dtype=computation_dtype)
+    delta: Array = lax.convert_element_type(stop - start, computation_dtype) / asarray(div, dtype=computation_dtype)
     iota_shape = [1,] * len(bounds_shape)
     iota_shape[axis] = div
     # This approach recovers the endpoints with float32 arithmetic,
     # but can lead to rounding errors for integer outputs.
     real_dtype = dtypes.finfo(computation_dtype).dtype
-    step = lax.iota(real_dtype, div).reshape(iota_shape) / jax.numpy.array(div, real_dtype)
+    step = lax.iota(real_dtype, div).reshape(iota_shape) / asarray(div, real_dtype)
     step = step.astype(computation_dtype)
     out = (broadcast_start.reshape(bounds_shape) * (1 - step) +
       broadcast_stop.reshape(bounds_shape) * step)
@@ -545,7 +549,7 @@ def _linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
                             canonicalize_axis(axis, out.ndim))
 
   elif num == 1:
-    delta = jax.numpy.asarray(np.nan if endpoint else stop - start, dtype=computation_dtype)
+    delta = asarray(np.nan if endpoint else stop - start, dtype=computation_dtype)
     out = broadcast_start.reshape(bounds_shape)
   else:  # num == 0 degenerate case, match numpy behavior
     empty_shape = list(lax.broadcast_shapes(np.shape(start), np.shape(stop)))
@@ -557,7 +561,7 @@ def _linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
     out = lax.floor(out)
 
   sharding = util.canonicalize_device_to_sharding(device)
-  result = lax_internal._convert_element_type(out, dtype, sharding=sharding)
+  result = lax._convert_element_type(out, dtype, sharding=sharding)
   return (result, delta) if retstep else result
 
 

jax/_src/numpy/array_methods.py

@@ -29,9 +29,8 @@
 from typing import Any, Callable, Sequence
 
 import numpy as np
-import jax
+
 from jax import lax
-from jax.sharding import Sharding
 from jax._src import api
 from jax._src import core
 from jax._src import dtypes
@@ -44,6 +43,7 @@
 from jax._src.numpy import lax_numpy
 from jax._src.numpy import tensor_contractions
 from jax._src.pjit import PartitionSpec
+from jax._src.sharding import Sharding
 from jax._src.sharding_impls import canonicalize_sharding, NamedSharding
 from jax._src.numpy import reductions
 from jax._src.numpy import ufuncs
@@ -612,12 +612,13 @@ def _deepcopy(self: Array, memo: Any) -> Array:
 
 def __array_module__(self, types):
   if all(issubclass(t, _HANDLED_ARRAY_TYPES) for t in types):
+    import jax.numpy  # pytype: disable=import-error
     return jax.numpy
   else:
     return NotImplemented
 
 
-@partial(jax.jit, static_argnums=(1,2,3))
+@partial(api.jit, static_argnums=(1,2,3))
 def _multi_slice(self: Array,
                  start_indices: tuple[tuple[int, ...]],
                  limit_indices: tuple[tuple[int, ...]],
@@ -637,7 +638,7 @@ def _multi_slice(self: Array,
 
 # The next two functions are related to iter(array), implemented here to
 # avoid circular imports.
-@jax.jit
+@api.jit
 def _unstack(x: Array) -> list[Array]:
   dims = (0,)
   return [lax.squeeze(t, dims) for t in lax.split(x, (1,) * x.shape[0])]
@@ -776,7 +777,7 @@ def __repr__(self) -> str:
     return f"_IndexUpdateRef({self.array!r}, {self.index!r})"
 
   def get(self, *, indices_are_sorted: bool = False, unique_indices: bool = False,
-          mode: str | jax.lax.GatherScatterMode | None = None,
+          mode: str | lax.GatherScatterMode | None = None,
           fill_value: ArrayLike | None = None, out_sharding: Sharding | None = None):
     """Equivalent to ``x[idx]``.
 
@@ -798,7 +799,7 @@ def get(self, *, indices_are_sorted: bool = False, unique_indices: bool = False,
 
   def set(self, values: ArrayLike, *, indices_are_sorted: bool = False,
           unique_indices: bool = False,
-          mode: str | jax.lax.GatherScatterMode | None = None) -> None:
+          mode: str | lax.GatherScatterMode | None = None) -> None:
     """Pure equivalent of ``x[idx] = y``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -816,7 +817,7 @@ def set(self, values: ArrayLike, *, indices_are_sorted: bool = False,
 
   def apply(self, func: Callable[[ArrayLike], Array], *,
             indices_are_sorted: bool = False, unique_indices: bool = False,
-            mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+            mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``func.at(x, idx)`` for a unary ufunc ``func``.
 
     Returns the value of ``x`` that would result from applying the unary
@@ -840,7 +841,7 @@ def _scatter_apply(x, indices, y, dims, **kwargs):
 
   def add(self, values: ArrayLike, *,
           indices_are_sorted: bool = False, unique_indices: bool = False,
-          mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+          mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] += y``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -855,7 +856,7 @@ def add(self, values: ArrayLike, *,
 
   def subtract(self, values: ArrayLike, *,
                indices_are_sorted: bool = False, unique_indices: bool = False,
-               mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+               mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] -= y``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -870,7 +871,7 @@ def subtract(self, values: ArrayLike, *,
 
   def multiply(self, values: ArrayLike, *,
                indices_are_sorted: bool = False, unique_indices: bool = False,
-               mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+               mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] *= y``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -887,7 +888,7 @@ def multiply(self, values: ArrayLike, *,
 
   def divide(self, values: ArrayLike, *,
              indices_are_sorted: bool = False, unique_indices: bool = False,
-             mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+             mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] /= y``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -904,7 +905,7 @@ def divide(self, values: ArrayLike, *,
 
   def power(self, values: ArrayLike, *,
             indices_are_sorted: bool = False, unique_indices: bool = False,
-            mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+            mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] **= y``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -921,7 +922,7 @@ def power(self, values: ArrayLike, *,
 
   def min(self, values: ArrayLike, *,
           indices_are_sorted: bool = False, unique_indices: bool = False,
-          mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+          mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] = minimum(x[idx], y)``.
 
     Returns the value of ``x`` that would result from the NumPy-style
@@ -937,7 +938,7 @@ def min(self, values: ArrayLike, *,
 
   def max(self, values: ArrayLike, *,
           indices_are_sorted: bool = False, unique_indices: bool = False,
-          mode: str | jax.lax.GatherScatterMode | None = None) -> Array:
+          mode: str | lax.GatherScatterMode | None = None) -> Array:
     """Pure equivalent of ``x[idx] = maximum(x[idx], y)``.
 
     Returns the value of ``x`` that would result from the NumPy-style

jax/_src/numpy/error.py

@@ -15,9 +15,11 @@
 import contextlib
 from typing import Literal, Sequence
 
-import jax
+import numpy as np
+
 from jax._src import config
-from jax._src.typing import ArrayLike
+from jax._src import dtypes
+from jax._src.typing import Array, ArrayLike
 
 Category = Literal["nan", "divide", "oob"]
 
@@ -40,7 +42,7 @@ def _is_category_disabled(
 
 
 def _set_error_if_with_category(
-    pred: jax.Array,
+    pred: Array,
     /,
     msg: str,
     category: Category | None = None,
@@ -65,7 +67,7 @@ def _set_error_if_with_category(
   error_check_lib.set_error_if(pred, msg)
 
 
-def _set_error_if_nan(pred: jax.Array, /):
+def _set_error_if_nan(pred: Array, /):
   """Set the internal error state if any element of `pred` is `NaN`.
 
   This function is disabled if the `jax_error_checking_behavior_nan` flag is
@@ -74,17 +76,17 @@ def _set_error_if_nan(pred: jax.Array, /):
   if config.error_checking_behavior_nan.value == "ignore":
     return
 
-  # TODO(mattjj): fix the circular import issue.
-  import jax.numpy as jnp
-  if not jnp.issubdtype(pred.dtype, jnp.floating):  # only check floats
+  if not dtypes.issubdtype(pred.dtype, np.floating):  # only check floats
     return
 
   # TODO(mattjj): fix the circular import issue.
   from jax._src import error_check as error_check_lib
+  import jax.numpy as jnp
+
   error_check_lib.set_error_if(jnp.isnan(pred), "NaN encountered")
 
 
-def _set_error_if_divide_by_zero(pred: jax.Array, /):
+def _set_error_if_divide_by_zero(pred: Array, /):
   """Set the internal error state if any element of `pred` is zero.
 
   This function is intended for checking if the denominator of a division is