Repository: google/jax
Pull request #20787 (2024-04-16 20:18:44)

Import rich lazily
This ensures that the timing of `import jax` is not affected by `rich` being installed.

See also #20778.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a30f333d9..331554c67 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,9 +9,8 @@ Remember to align the itemized text with the first line of an item within a list
 ## jax 0.4.27
 
 * New Functionality
-  * Added {func}`jax.numpy.unstack` and {func}`jax.numpy.cumulative_sum`,
-    following their addition in the array API 2023 standard, soon to be
-    adopted by NumPy.
+  * Added {func}`jax.numpy.unstack`, following the addition of this function in
+    the array API 2023 standard, soon to be adopted by NumPy.
 
 * Changes
   * {func}`jax.pure_callback` and {func}`jax.experimental.io_callback`
diff --git a/docs/jax.numpy.rst b/docs/jax.numpy.rst
index 5ccf043d2..d866df52d 100644
--- a/docs/jax.numpy.rst
+++ b/docs/jax.numpy.rst
@@ -138,7 +138,6 @@ namespace; they are listed below.
     csingle
     cumprod
     cumsum
-    cumulative_sum
     deg2rad
     degrees
     delete
diff --git a/jax/_src/compilation_cache_interface.py b/jax/_src/compilation_cache_interface.py
index 701bd6bf3..95d557c55 100644
--- a/jax/_src/compilation_cache_interface.py
+++ b/jax/_src/compilation_cache_interface.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 
 from abc import abstractmethod
-import pathlib
 
+from jax._src import path as pathlib
 from jax._src import util
 
 
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 84fc677c9..e923d1e5f 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import importlib.util
 from collections.abc import Sequence
 import functools
 import string
@@ -44,20 +45,6 @@ from jax._src.lib.mlir.dialects import hlo
 from jax._src.sharding import Sharding
 from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
 
-# pytype: disable=import-error
-try:
-  import rich
-  import rich.align
-  import rich.box
-  import rich.console
-  import rich.padding
-  import rich.style
-  import rich.table
-  RICH_ENABLED = True
-except:
-  RICH_ENABLED = False
-# pytype: enable=import-error
-
 class DebugEffect(effects.Effect):
   __str__ = lambda self: "Debug"
 debug_effect = DebugEffect()
@@ -440,8 +427,19 @@ def visualize_sharding(shape: Sequence[int], sharding: Sharding, *,
                        min_width: int = 9, max_width: int = 80,
                        color_map: ColorMap | None = None):
   """Visualizes a ``Sharding`` using ``rich``."""
-  if not RICH_ENABLED:
+  if not importlib.util.find_spec("rich"):
     raise ValueError("`visualize_sharding` requires `rich` to be installed.")
+
+  # These imports are local so that they don't affect JAX import times.
+  # pytype: disable=import-error
+  import rich.align
+  import rich.console
+  import rich.box
+  import rich.padding
+  import rich.style
+  import rich.table
+  # pytype: enable=import-error
+
   if len(shape) > 2 or len(shape) < 1:
     raise ValueError(
         "`visualize_sharding` only works for shapes with 1 and 2 dimensions.")
diff --git a/jax/_src/gfile_cache.py b/jax/_src/gfile_cache.py
index 95505e858..301f61cc6 100644
--- a/jax/_src/gfile_cache.py
+++ b/jax/_src/gfile_cache.py
@@ -21,7 +21,7 @@ class GFileCache(CacheInterface):
 
   def __init__(self, path: str):
     """Sets up a cache at 'path'. Cached values may already be present."""
-    self._path = pathlib.get_path(path)
+    self._path = pathlib.Path(path)
     self._path.mkdir(parents=True, exist_ok=True)
 
   def get(self, key: str):
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index ad7ab2af3..b923d7185 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -464,7 +464,7 @@ def dump_module_to_file(module: ir.Module, stage_name: str) -> str | None:
 
   name = f"jax_ir{id}_{_make_string_safe_for_filename(module_name)}_{stage_name}.mlir"
 
-  out_dir = path.get_path(out_dir_name)
+  out_dir = path.Path(out_dir_name)
   out_dir.mkdir(parents=True, exist_ok=True)
   full_path = out_dir / name
   full_path.write_text(module_to_string(module))
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index 8c5594cb4..0e628fc5e 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -26,7 +26,7 @@ import numpy as np
 
 from jax import lax
 from jax._src import api
-from jax._src import core, config
+from jax._src import core
 from jax._src import dtypes
 from jax._src.numpy import ufuncs
 from jax._src.numpy.util import (
@@ -708,42 +708,6 @@ nancumsum = _make_cumulative_reduction(np.nancumsum, lax.cumsum,
 nancumprod = _make_cumulative_reduction(np.nancumprod, lax.cumprod,
                                         fill_nan=True, fill_value=1)
 
-@implements(getattr(np, 'cumulative_sum', None))
-def cumulative_sum(
-    x: ArrayLike, /, *, axis: int | None = None,
-    dtype: DTypeLike | None = None,
-    include_initial: bool = False) -> Array:
-  check_arraylike("cumulative_sum", x)
-  x = lax_internal.asarray(x)
-  if x.ndim == 0:
-    raise ValueError(
-      "The input must be non-scalar to take a cumulative sum, however a "
-      "scalar value or scalar array was given."
-    )
-  if axis is None:
-    axis = 0
-    if x.ndim > 1:
-      raise ValueError(
-        f"The input array has rank {x.ndim}, however axis was not set to an "
-        "explicit value. The axis argument is only optional for one-dimensional "
-        "arrays.")
-
-  axis = _canonicalize_axis(axis, x.ndim)
-  dtypes.check_user_dtype_supported(dtype)
-  kind = x.dtype.kind
-  if (dtype is None and kind in {'i', 'u'}
-      and x.dtype.itemsize*8 < int(config.default_dtype_bits.value)):
-    dtype = dtypes.canonicalize_dtype(dtypes._default_types[kind])
-  x = x.astype(dtype=dtype or x.dtype)
-  out = cumsum(x, axis=axis)
-  if include_initial:
-    zeros_shape = list(x.shape)
-    zeros_shape[axis] = 1
-    out = lax_internal.concatenate(
-      [lax_internal.full(zeros_shape, 0, dtype=out.dtype), out],
-      dimension=axis)
-  return out
-
 # Quantiles
 @implements(np.quantile, skip_params=['out', 'overwrite_input'])
 @partial(api.jit, static_argnames=('axis', 'overwrite_input', 'interpolation',
diff --git a/jax/_src/path.py b/jax/_src/path.py
index 37869a79a..1f92dd504 100644
--- a/jax/_src/path.py
+++ b/jax/_src/path.py
@@ -12,27 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
 import logging
 import pathlib
 
 logger = logging.getLogger(__name__)
 
-Path = pathlib.Path
-
-@functools.cache
-def get_path_implementation():
-  # We import etils lazily, because it can be slow to import.
-  try:
-    import etils.epath as epath
-    logger.debug("etils.epath found. Using etils.epath for file I/O.")
-    return epath.Path
-  except ImportError:
-    logger.debug("etils.epath was not found. Using pathlib for file I/O.")
-    return pathlib.Path
+try:
+  import etils.epath as epath
+except:
+  epath = None
 
 # If etils.epath (aka etils[epath] to pip) is present, we prefer it because it
 # can read and write to, e.g., GCS buckets. Otherwise we use the builtin
 # pathlib and can only read/write to the local filesystem.
-def get_path(*args) -> Path:
-  return get_path_implementation()(*args)
+if epath:
+  logger.debug("etils.epath found. Using etils.epath for file I/O.")
+  Path = epath.Path
+else:
+  logger.debug("etils.epath was not found. Using pathlib for file I/O.")
+  Path = pathlib.Path
diff --git a/jax/_src/xla_bridge.py b/jax/_src/xla_bridge.py
index 0c7ab97f9..339c9ea00 100644
--- a/jax/_src/xla_bridge.py
+++ b/jax/_src/xla_bridge.py
@@ -106,11 +106,9 @@ _CPU_ENABLE_GLOO_COLLECTIVES = config.DEFINE_bool(
     help="If True, enable cross-process collectives on CPU using Gloo.",
 )
 
-# TODO(yueshengys): turn default back to True after resolving memory increase
-# issue.
 _CPU_ENABLE_ASYNC_DISPATCH = config.DEFINE_bool(
     name="jax_cpu_enable_async_dispatch",
-    default=False,
+    default=True,
     help="Only applies to non-parallel computations. If False, run computations"
     "inline without async dispatch.",
 )
diff --git a/jax/experimental/array_api/__init__.py b/jax/experimental/array_api/__init__.py
index 7050a4e3c..dfb2ff988 100644
--- a/jax/experimental/array_api/__init__.py
+++ b/jax/experimental/array_api/__init__.py
@@ -204,7 +204,6 @@ from jax.experimental.array_api._sorting_functions import (
 )
 
 from jax.experimental.array_api._statistical_functions import (
-    cumulative_sum as cumulative_sum,
     max as max,
     mean as mean,
     min as min,
diff --git a/jax/experimental/array_api/_statistical_functions.py b/jax/experimental/array_api/_statistical_functions.py
index 141b80abf..2e1333317 100644
--- a/jax/experimental/array_api/_statistical_functions.py
+++ b/jax/experimental/array_api/_statistical_functions.py
@@ -18,10 +18,6 @@ from jax.experimental.array_api._data_type_functions import (
 )
 
 
-def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False):
-  """Calculates the cumulative sum of elements in the input array x."""
-  return jax.numpy.cumulative_sum(x, axis=axis, dtype=dtype, include_initial=include_initial)
-
 def max(x, /, *, axis=None, keepdims=False):
   """Calculates the maximum value of the input array x."""
   return jax.numpy.max(x, axis=axis, keepdims=keepdims)
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
index 1b9a990f3..0b8a8cdd0 100644
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -296,7 +296,6 @@ from jax._src.numpy.reductions import (
     count_nonzero as count_nonzero,
     cumsum as cumsum,
     cumprod as cumprod,
-    cumulative_sum as cumulative_sum,
     max as max,
     mean as mean,
     median as median,
diff --git a/jax/numpy/__init__.pyi b/jax/numpy/__init__.pyi
index 274063804..7c7e68a06 100644
--- a/jax/numpy/__init__.pyi
+++ b/jax/numpy/__init__.pyi
@@ -241,9 +241,6 @@ def cumprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
 cumproduct = cumprod
 def cumsum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ...) -> Array: ...
-def cumulative_sum(x: ArrayLike, /, *, axis: int | None = ...,
-                   dtype: DTypeLike | None = ...,
-                   include_initial: bool = ...) -> Array: ...
 
 def deg2rad(x: ArrayLike, /) -> Array: ...
 degrees = rad2deg


Test patch:
diff --git a/tests/array_api_test.py b/tests/array_api_test.py
index 5667c3459..91c64e749 100644
--- a/tests/array_api_test.py
+++ b/tests/array_api_test.py
@@ -68,7 +68,6 @@ MAIN_NAMESPACE = {
   'copysign',
   'cos',
   'cosh',
-  'cumulative_sum',
   'divide',
   'e',
   'empty',
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
index 73100352c..b72e47d9d 100644
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -770,64 +770,5 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
     self.assertAllClose(expected, actual, atol=0)
 
 
-  @jtu.sample_product(
-    [dict(shape=shape, axis=axis)
-      for shape in all_shapes
-      for axis in list(
-        range(-len(shape), len(shape))
-      ) + ([None] if len(shape) == 1 else [])],
-    dtype=all_dtypes + [None],
-    out_dtype=all_dtypes,
-    include_initial=[False, True],
-  )
-  @jtu.ignore_warning(category=NumpyComplexWarning)
-  @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises mixed type promotion
-  def testCumulativeSum(self, shape, axis, dtype, out_dtype, include_initial):
-    rng = jtu.rand_some_zero(self.rng())
-
-    def np_mock_op(x, axis=None, dtype=None, include_initial=False):
-      kind = x.dtype.kind
-      if (dtype is None and kind in {'i', 'u'}
-          and x.dtype.itemsize*8 < int(config.default_dtype_bits.value)):
-        dtype = dtypes.canonicalize_dtype(dtypes._default_types[kind])
-      axis = axis or 0
-      x = x.astype(dtype=dtype or x.dtype)
-      out = jnp.cumsum(x, axis=axis)
-      if include_initial:
-        zeros_shape = list(x.shape)
-        zeros_shape[axis] = 1
-        out = jnp.concat([jnp.zeros(zeros_shape, dtype=out.dtype), out], axis=axis)
-      return out
-
-
-    # We currently "cheat" to ensure we have JAX arrays, not NumPy arrays as
-    # input because we rely on JAX-specific casting behavior
-    args_maker = lambda: [jnp.array(rng(shape, dtype))]
-    np_op = getattr(np, "cumulative_sum", np_mock_op)
-    kwargs = dict(axis=axis, dtype=out_dtype, include_initial=include_initial)
-
-    np_fun = lambda x: np_op(x, **kwargs)
-    jnp_fun = lambda x: jnp.cumulative_sum(x, **kwargs)
-    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
-    self._CompileAndCheck(jnp_fun, args_maker)
-
-
-  @jtu.sample_product(
-      shape=filter(lambda x: len(x) != 1, all_shapes), dtype=all_dtypes,
-      include_initial=[False, True])
-  def testCumulativeSumErrors(self, shape, dtype, include_initial):
-    rng = jtu.rand_some_zero(self.rng())
-    x = rng(shape, dtype)
-    rank = jnp.asarray(x).ndim
-    if rank == 0:
-      msg = r"The input must be non-scalar to take"
-      with self.assertRaisesRegex(ValueError, msg):
-        jnp.cumulative_sum(x, include_initial=include_initial)
-    elif rank > 1:
-      msg = r"The input array has rank \d*, however"
-      with self.assertRaisesRegex(ValueError, msg):
-        jnp.cumulative_sum(x, include_initial=include_initial)
-
-
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())