Merge pull request #136 from ROCm/ci-upstream-sync-16_1

CI: 11/12/24 upstream sync

Author: charleshofer

.github/workflows/ci-build.yaml

@@ -35,11 +35,11 @@ jobs:
         with:
           python-version: 3.11
       - run: python -m pip install pre-commit
-      - uses: actions/cache@3624ceb22c1c5a301c8db4169662070a689d9ea8 # v4.1.1
+      - uses: actions/cache@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.1.2
         with:
           path: ~/.cache/pre-commit
           key: pre-commit-${{ env.pythonLocation }}-${{ hashFiles('.pre-commit-config.yaml', 'setup.py') }}
-      - run: pre-commit run --show-diff-on-failure --color=always
+      - run: pre-commit run --show-diff-on-failure --color=always --all-files
 
   build:
     name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ubuntu-20.04, x64=${{ matrix.enable-x64}})"

CHANGELOG.md

@@ -25,6 +25,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     for information on migrating to the new API.
   * The `initial` argument to {func}`jax.nn.softmax` and {func}`jax.nn.log_softmax`
     has been removed, after being deprecated in v0.4.27.
+  * Calling `np.asarray` on typed PRNG keys (i.e. keys produced by :func:`jax.random.key`)
+    now raises an error. Previously, this returned a scalar object array.
   * The following deprecated methods and functions in {mod}`jax.export` have
     been removed:
       * `jax.export.DisabledSafetyCheck.shape_assertions`: it had no effect

docs/about.md

@@ -0,0 +1,123 @@
+(about-the-project)=
+
+# About the project
+
+The JAX project is led by the JAX core team. We develop in the open,
+and welcome open-source contributions from across the community. We
+frequently see contributions from [Google
+DeepMind](https://deepmind.google/), Alphabet more broadly,
+[NVIDIA](https://docs.nvidia.com/deeplearning/frameworks/jax-release-notes/overview.html),
+and elsewhere.
+
+At the heart of the project is the [JAX
+core](http://github.com/google/jax) library, which focuses on the
+fundamentals of machine learning and numerical computing, at scale.
+
+When [developing](#development) the core, we want to maintain agility
+and a focused scope, so we lean heavily on a surrounding [modular
+technology stack](#components). First, we design the `jax` module
+to be
+[composable](https://github.com/jax-ml/jax?tab=readme-ov-file#transformations)
+and
+[extensible](https://jax.readthedocs.io/en/latest/jax.extend.html), so
+that a wide variety of domain-specific libraries can thrive outside of
+it in a decentralized manner. Second, we lean heavily on a modular
+backend stack (compiler and runtime) to target different
+accelerators. Whether you are [writing a new domain-specific library
+built with JAX](#upstack), or looking to [support
+new hardware](#downstack), you can often
+contribute these with *minimal to no modifications* to the JAX core
+codebase.
+
+Many of JAX's core contributors have roots in open-source software and
+in research, in fields spanning computer science and the natural
+sciences. We strive to continuously enable the cutting edge of machine
+learning and numerical computing---across all compute platforms and
+accelerators---and to discover the truths of array programming at
+scale.
+
+(development)=
+## Open development
+
+JAX's day-to-day development takes place in the open on GitHub, using
+pull requests, the issue tracker, discussions, and [JAX Enhancement
+Proposals
+(JEPs)](https://jax.readthedocs.io/en/latest/jep/index.html). Reading
+and participating in these is a good way to get involved. We also
+maintain [developer
+notes](https://jax.readthedocs.io/en/latest/contributor_guide.html)
+that cover JAX's internal design.
+
+The JAX core team determines whether to accept changes and
+enhancements. Maintaining a simple decision-making structure currently
+helps us develop at the speed of the research frontier. Open
+development is a core value of ours, and we may adapt to a more
+intricate decision structure over time (e.g. with designated area
+owners) if/when it becomes useful to do so.
+
+For more see [contributing to
+JAX](https://jax.readthedocs.io/en/latest/contributing.html).
+
+(components)=
+## A modular stack
+
+To enable (a) a growing community of users across numerical domains,
+and (b) an advancing hardware landscape, we lean heavily on
+**modularity**.
+
+(upstack)=
+### Libraries built on JAX
+
+While the JAX core library focuses on the fundamentals, we want to
+encourage domain-specific libraries and tools to be built on top of
+JAX. Indeed, [many
+libraries](https://jax.readthedocs.io/en/latest/#ecosystem) have
+emerged around JAX to offer higher-level features and extensions.
+
+How do we encourage such decentralized development? We guide it with
+several technical choices. First, JAX's main API focuses on basic
+building blocks (e.g. numerical primitives, NumPy operations, arrays,
+and transformations), encouraging auxiliary libraries to develop
+utilities as needed for their domain. In addition, JAX exposes a
+handful of more advanced APIs for
+[customization](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html)
+and
+[extensibility](https://jax.readthedocs.io/en/latest/jax.extend.html). Libraries
+can [lean on these
+APIs](https://jax.readthedocs.io/en/latest/building_on_jax.html) in
+order to use JAX as an internal means of implementation, to integrate
+more with its transformations like autodiff, and more.
+
+Projects across the JAX ecosystem are developed in a distributed and
+often open fashion. They are not governed by the JAX core team, even
+though sometimes team members contribute to them or maintain contact
+with their developers.
+
+(downstack)=
+### A pluggable backend
+
+We want JAX to run on CPUs, GPUs, TPUs, and other hardware platforms
+as they emerge. To encourage unhindered support of JAX on new
+platforms, the JAX core emphasizes modularity in its backend too.
+
+To manage hardware devices and memory, and for compilation to such
+devices, JAX calls out to the open [XLA
+compiler](https://openxla.org/) and the [PJRT
+runtime](https://github.com/openxla/xla/tree/main/xla/pjrt/c#pjrt---uniform-device-api). Both
+of these are projects external to the JAX core, governed and
+maintained by OpenXLA (again, with frequent contributions from and
+discussion with the JAX core developers).
+
+XLA aims for interoperability across accelerators (e.g. by ingesting
+[StableHLO](https://openxla.org/stablehlo) as input) and PJRT offers
+extensibility through a plug-in device API. Adding support for new
+devices is done by implementing a backend lowering for XLA, and
+implementing a plug-in device API defined by PJRT. If you're looking
+to contribute to compilation, or to supporting new hardware, we
+encourage you to contribute at the XLA and PJRT layers.
+
+These open system components allow third parties to support JAX on new
+accelerator platforms, *without requiring changes in the JAX
+core*. There are several plug-ins in development today. For example, a
+team at Apple is working on a PJRT plug-in to get [JAX running on
+Apple Metal](https://developer.apple.com/metal/jax/).

docs/index.rst

@@ -142,6 +142,7 @@ maintains an up-to-date list.
    extensions
    notes
    jax
+   about
 
 
 .. toctree::

jax/_src/api.py

@@ -2445,12 +2445,8 @@ def _device_get(x):
 
   # Extended dtypes dispatch via their device_get rule.
   if isinstance(x, basearray.Array) and dtypes.issubdtype(x.dtype, dtypes.extended):
-    try:
-      to_device = x.dtype._rules.device_get
-    except AttributeError:
-      pass
-    else:
-      return to_device(x)
+    bufs, tree = tree_util.dispatch_registry.flatten(x)
+    return tree.unflatten(device_get(bufs))
 
   # Other types dispatch via their __array__ method.
   try:

jax/_src/array.py

@@ -1173,7 +1173,7 @@ def _array_shard_arg(xs, shardings, layouts, copy_semantics):
     copy_outs = xc.batched_copy_array_to_devices_with_sharding(
         batch_xs, batch_devs, batch_shardings, batch_cs)
   else:
-    copy_outs = xc.batched_copy_array_to_devices_with_sharding(  # type: ignore
+    copy_outs = xc.batched_copy_array_to_devices_with_sharding(  # pytype: disable=missing-parameter
         batch_xs, batch_devs, batch_shardings)
   for i, copy_out in safe_zip(batch_indices, copy_outs):
     assert results[i] is None

jax/_src/compilation_cache.py

@@ -84,6 +84,8 @@ def is_cache_used(backend: xla_client.Client) -> bool:
         _cache_used = True
       return _cache_used
 
+  return False
+
 
 def get_file_cache(path: str) -> tuple[CacheInterface, str] | None:
   """Returns the file cache and the path to the cache."""

jax/_src/config.py

@@ -22,7 +22,7 @@
 import os
 import sys
 import threading
-from typing import Any, Generic, NamedTuple, NoReturn, Optional, Protocol, TypeVar, cast, TYPE_CHECKING
+from typing import Any, Generic, NamedTuple, NoReturn, Optional, Protocol, TypeVar, cast
 
 from jax._src import lib
 from jax._src.lib import guard_lib
@@ -371,7 +371,7 @@ class _Unset: pass
 
   _thread_local_state = threading.local()
 
-  class State(Generic[_T]):
+  class State(Generic[_T]):  # type: ignore[no-redef]
 
     __slots__ = (
         '_name', '_value', '_update_thread_local_hook', '_update_global_hook',

jax/_src/core.py

@@ -892,6 +892,11 @@ def unsafe_buffer_pointer(self):
 aval_property = namedtuple("aval_property", ["fget"])
 aval_method = namedtuple("aval_method", ["fun"])
 
+def check_eval_args(args):
+  for arg in args:
+    if isinstance(arg, Tracer):
+      raise escaped_tracer_error(arg)
+
 class EvalTrace(Trace):
 
   def process_primitive(self, primitive, args, params):
@@ -902,12 +907,11 @@ def process_primitive(self, primitive, args, params):
     else:
       # TODO(dougalm): delete. this shouldn't be necessary
       args = map(full_lower, args)
-      for arg in args:
-        if isinstance(arg, Tracer):
-          if config.data_dependent_tracing_fallback.value:
+      if config.data_dependent_tracing_fallback.value:
+        for arg in args:
+          if isinstance(arg, Tracer):
             return primitive.bind_with_trace(arg._trace, args, params)
-          else:
-            raise escaped_tracer_error(arg)
+      check_eval_args(args)
       return primitive.impl(*args, **params)
 
   def process_call(self, primitive, f, tracers, params):
@@ -955,6 +959,7 @@ def __eq__(self, other):
 @dataclass(frozen=True)
 class AxisEnv:
   axis_sizes : dict[AxisName, int]
+  spmd_axis_names : set[AxisName]
 
   def axis_size(self, axis_name):
     if axis_name not in self.axis_sizes:
@@ -971,20 +976,24 @@ def axis_names(self):
   def pop_pure(self, axis_name):
     new_sizes = self.axis_sizes.copy()
     new_sizes.pop(axis_name)
-    return AxisEnv(new_sizes)
+    return AxisEnv(new_sizes, self.spmd_axis_names)
 
   def extend_pure(self, name_size_pairs):
     new_sizes = self.axis_sizes.copy()
     new_sizes.update((name, size) for name, size in name_size_pairs
                     if name is not no_axis_name)
-    return AxisEnv(new_sizes)
+    return AxisEnv(new_sizes, self.spmd_axis_names)
+
+  def add_spmd_axis_names(self, axis_names):
+    new_spmd_axis_names = self.spmd_axis_names | set(axis_names)
+    return AxisEnv(self.axis_sizes, new_spmd_axis_names)
 
   def as_hashable_key(self):
     return tuple((name, size) for (name, size) in self.axis_sizes.items()
                  if name is not no_axis_name)
 
 eval_trace = EvalTrace()
-top_axis_env = AxisEnv({})
+top_axis_env = AxisEnv({}, set())
 
 class TracingContext(threading.local):
   trace: Trace | None
@@ -1045,6 +1054,16 @@ def extend_axis_env_nd(name_size_pairs : Iterable[tuple[AxisName, int]]):
   finally:
     trace_ctx.set_axis_env(prev)
 
+@contextmanager
+def add_spmd_axis_names(axis_names: AxisName | None):
+  prev = trace_ctx.axis_env
+  try:
+    if axis_names is not None:
+      trace_ctx.set_axis_env(prev.add_spmd_axis_names(axis_names))
+    yield
+  finally:
+    trace_ctx.set_axis_env(prev)
+
 def get_axis_env():
   return trace_ctx.axis_env
 
@@ -2092,33 +2111,6 @@ def get_bind_params(self, params):
 closed_call_p.def_effectful_abstract_eval(
     lambda *_, call_jaxpr: (call_jaxpr.out_avals, call_jaxpr.effects))
 
-
-outfeed_primitives: set[Primitive] = set()
-def jaxpr_uses_outfeed(jaxpr: Jaxpr) -> bool:
-  """Finds if there are outfeed primitives anywhere inside a Jaxpr."""
-  return any(primitive_uses_outfeed(eqn.primitive, eqn.params)
-             for eqn in jaxpr.eqns)
-
-def _param_uses_outfeed(param):
-  if type(param) is Jaxpr:
-    if jaxpr_uses_outfeed(param):
-      return True
-  elif type(param) is ClosedJaxpr:
-    if jaxpr_uses_outfeed(param.jaxpr):
-      return True
-  return False
-
-def primitive_uses_outfeed(prim: Primitive, params: dict) -> bool:
-  if prim in outfeed_primitives:
-    return True
-  for param in params.values():
-    if isinstance(param, tuple):
-      if any(unsafe_map(_param_uses_outfeed, param)):
-        return True
-    elif _param_uses_outfeed(param):
-      return True
-  return False
-
 # ------------------- Map -------------------
 
 class MapPrimitive(Primitive):

jax/_src/dispatch.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 import atexit
-from collections.abc import Callable, Sequence
+from collections.abc import Sequence
 import contextlib
 import dataclasses
 import enum
@@ -278,17 +278,6 @@ def _is_bint_axis_size(d: core.AxisSize) -> bool:
   return False
 
 
-# We can optionally set a Jaxpr rewriter that can be applied just before
-# compilation. This mechanism is used for compiling id_tap, we can
-# remove it once we bring the id_tap implementation into the core.
-outfeed_rewriter: Callable[[core.Jaxpr], core.Jaxpr] | None = None
-def apply_outfeed_rewriter(jaxpr: core.Jaxpr) -> core.Jaxpr:
-  if outfeed_rewriter is not None:
-    return outfeed_rewriter(jaxpr)
-  else:
-    return jaxpr
-
-
 def check_arg(arg: Any):
   if not (isinstance(arg, core.Tracer) or core.valid_jaxtype(arg)):
     raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid "