Add ResultHandler.wrap which allows us to avoid exploding PRNG keys.

PiperOrigin-RevId: 842010889

Author: pschuh

jax/_src/array.py

@@ -39,6 +39,7 @@
 from jax._src.interpreters import pxla
 from jax._src.layout import AutoLayout, Format, Layout
 from jax._src.lib import _jax
+from jax._src.lib import jaxlib_extension_version
 from jax._src.lib import xla_client as xc
 from jax._src.mesh import empty_concrete_mesh
 from jax._src.sharding import Sharding
@@ -1285,7 +1286,14 @@ def _array_shard_arg(xs, shardings, layouts, copy_semantics):
 
 def _array_global_result_handler(global_aval, out_sharding, committed):
   if global_aval.dtype == dtypes.float0:
-    return lambda _: np.zeros(global_aval.shape, dtypes.float0)
+    def handler(xs):
+      return np.zeros(global_aval.shape, dtypes.float0)
+    if jaxlib_extension_version >= 390:
+      phys_aval = core.physical_aval(global_aval)
+      return xc.array_result_handler(phys_aval, out_sharding, committed=committed,
+                                     _skip_checks=True).wrap(handler)
+    else:
+      return handler
   if dtypes.issubdtype(global_aval.dtype, dtypes.extended):
     return global_aval.dtype._rules.global_sharded_result_handler(
         global_aval, out_sharding, committed)
@@ -1297,7 +1305,14 @@ def _array_global_result_handler(global_aval, out_sharding, committed):
 # Only used for Arrays that come out of pmap.
 def _array_local_result_handler(aval, sharding, indices):
   if aval.dtype == dtypes.float0:
-    return lambda _: np.zeros(aval.shape, dtypes.float0)
+    def handler(xs):
+      return np.zeros(aval.shape, dtypes.float0)
+    if jaxlib_extension_version >= 390:
+      phys_aval = core.physical_aval(aval)
+      return xc.array_result_handler(phys_aval, sharding, committed=True,
+                                     _skip_checks=True).wrap(handler)
+    else:
+      return handler
   if dtypes.issubdtype(aval.dtype, dtypes.extended):
     return aval.dtype._rules.local_sharded_result_handler(
         aval, sharding, indices)
@@ -1326,9 +1341,13 @@ def _token_shard_arg(xs, shardings, layouts, copy_semantics):
 def _token_global_result_handler(global_aval, out_sharding, committed):
   array_handler = _array_global_result_handler(
       core.get_token_aval(), out_sharding, committed)
-
-  def wrapper(*args, **kwargs):
-    out_buf = array_handler(*args, **kwargs)
-    return core.Token(out_buf)
-  return wrapper
+  if jaxlib_extension_version >= 390:
+    def wrapper(array):
+      return core.Token(array)
+    return array_handler.wrap(wrapper)  # type: ignore
+  else:
+    def old_wrapper(*args, **kwargs):
+      out_buf = array_handler(*args, **kwargs)
+      return core.Token(out_buf)
+    return old_wrapper
 pxla.global_result_handlers[core.AbstractToken] = _token_global_result_handler

jax/_src/interpreters/pxla.py

@@ -302,7 +302,7 @@ def local_aval_to_result_handler(
     raise TypeError(
         f"No pxla_result_handler for type: {type(aval)}") from err
 
-PxlaResultHandler = Callable[..., Callable[[Any], Any]]
+PxlaResultHandler = Callable[..., xc._xla.ResultHandler]
 local_result_handlers: dict[type[core.AbstractValue], PxlaResultHandler] = {}
 
 

jax/_src/lax/lax.py

@@ -63,6 +63,7 @@
 from jax._src.lax.utils import (
   input_dtype, dtype_to_string, standard_multi_result_abstract_eval,
   standard_primitive)
+from jax._src.lib import jaxlib_extension_version
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import chlo
 from jax._src.lib.mlir.dialects import hlo
@@ -9246,10 +9247,14 @@ def global_sharded_result_handler(aval, out_sharding, committed):
     else:
       phys_sharding = out_sharding
     phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed)
-
-    def handler(bufs):
-      return core.DArray(aval, phys_handler(bufs))
-    return handler
+    if jaxlib_extension_version >= 390:
+      def handler(arr):
+        return core.DArray(aval, arr)
+      return phys_handler.wrap(handler)
+    else:
+      def handler(bufs):
+        return core.DArray(aval, phys_handler(bufs))
+      return handler
 
 
 core.bint._rules = BIntRules

jax/_src/prng.py

@@ -42,6 +42,7 @@
 from jax._src.lax import control_flow as lax_control_flow
 from jax._src.lax import lax
 from jax._src.lax import slicing as lax_slicing
+from jax._src.lib import jaxlib_extension_version
 from jax._src.lib import gpu_prng
 from jax._src.lib import xla_client as xc
 from jax._src.lib.mlir import ir
@@ -402,10 +403,16 @@ def local_sharded_result_handler(aval, sharding, indices):
     phys_handler = phys_handler_maker(phys_aval, phys_sharding, phys_indices)
 
     # set up a handler that calls the physical one and wraps back up
-    def handler(bufs):
-      return PRNGKeyArray(aval.dtype._impl, phys_handler(bufs))
+    if jaxlib_extension_version >= 390:
+      def handler(arr):
+        return PRNGKeyArray(aval.dtype._impl, arr)
 
-    return handler
+      return phys_handler.wrap(handler)
+    else:
+      def handler(bufs):
+        return PRNGKeyArray(aval.dtype._impl, phys_handler(bufs))
+
+      return handler
 
   @staticmethod
   def global_sharded_result_handler(aval, out_sharding, committed):
@@ -414,9 +421,14 @@ def global_sharded_result_handler(aval, out_sharding, committed):
 
     phys_sharding = physical_sharding(aval, out_sharding)
     phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed)
-    def handler(bufs):
-      return PRNGKeyArray(aval.dtype._impl, phys_handler(bufs))
-    return handler
+    if jaxlib_extension_version >= 390:
+      def handler(bufs):
+        return PRNGKeyArray(aval.dtype._impl, bufs)
+      return phys_handler.wrap(handler)
+    else:
+      def handler(bufs):
+        return PRNGKeyArray(aval.dtype._impl, phys_handler(bufs))
+      return handler
 
   @staticmethod
   def make_sharded_array(aval, sharding, arrays, committed):

jaxlib/_jax/__init__.pyi

@@ -1029,6 +1029,7 @@ def array_result_handler(
 
 class ResultHandler:
   def __call__(self, arg: Array | Sequence[Array], /) -> Array: ...
+  def wrap(self, wrapper: Callable) -> Any: ...
 
 class DeviceList:
   def __init__(self, arg: tuple[Device, ...], /) -> None: ...

jaxlib/py_array.cc

@@ -584,18 +584,21 @@ PyArray PyArray::MakeFromIfrtArrayAndSharding(nb_class_ptr<PyClient> py_client,
                  std::move(ifrt_array), committed, skip_checks);
 }
 
-PyArrayResultHandler::PyArrayResultHandler(nb::object aval, nb::object sharding,
-                                           bool committed, bool skip_checks)
+PyArrayResultHandler::PyArrayResultHandler(
+    nb::object aval, nb::object sharding, bool committed, bool skip_checks,
+    std::vector<nanobind::callable> wrappers)
     : aval_(std::move(aval)),
       sharding_(std::move(sharding)),
       committed_(committed),
-      skip_checks_(skip_checks) {
+      skip_checks_(skip_checks),
+      wrappers_(std::move(wrappers)) {
   weak_type_ = nb::cast<bool>(aval_.attr("weak_type"));
   dtype_ = nb::cast<xla::nb_dtype>(aval_.attr("dtype"));
   shape_ = nb::cast<std::vector<int64_t>>(aval_.attr("shape"));
 }
 
-PyArray PyArrayResultHandler::Call(absl::Span<const PyArray> py_arrays) const {
+nanobind::object PyArrayResultHandler::Call(
+    absl::Span<const PyArray> py_arrays) const {
   auto py_device_list = GetPyDeviceList(sharding_);
   if (!py_device_list.ok()) {
     throw nb::value_error(
@@ -610,15 +613,20 @@ PyArray PyArrayResultHandler::Call(absl::Span<const PyArray> py_arrays) const {
               xla::Future<>());
 }
 
-PyArray PyArrayResultHandler::Call(nb_class_ptr<PyClient> py_client,
-                                   ifrt::ArrayRef ifrt_array,
-                                   xla::Future<> result_status) const {
-  return PyArray(aval_, weak_type_, dtype_, shape_, sharding_,
-                 std::move(py_client), std::move(ifrt_array), committed_,
-                 skip_checks_, std::move(result_status));
+nanobind::object PyArrayResultHandler::Call(nb_class_ptr<PyClient> py_client,
+                                            ifrt::ArrayRef ifrt_array,
+                                            xla::Future<> result_status) const {
+  nanobind::object result =
+      PyArray(aval_, weak_type_, dtype_, shape_, sharding_,
+              std::move(py_client), std::move(ifrt_array), committed_,
+              skip_checks_, std::move(result_status));
+  for (auto& cb : wrappers_) {
+    result = cb(std::move(result));
+  }
+  return result;
 }
 
-PyArray PyArrayResultHandler::Call(PyArray py_array) const {
+nanobind::object PyArrayResultHandler::Call(PyArray py_array) const {
   return Call(py_array.py_client(), tsl::FormRef(py_array.ifrt_array()),
               xla::Future<>());
 }
@@ -2364,7 +2372,14 @@ absl::Status PyArray::Register(nb::module_& m) {
                     .c_str());
           },
           nb::sig(
-              "def __call__(self, arg: Array | Sequence[Array], /) -> Array"));
+              "def __call__(self, arg: Array | Sequence[Array], /) -> Array"))
+      .def("wrap", [](const PyArrayResultHandler& self, nb::callable wrapper) {
+        auto wrappers = self.wrappers();
+        wrappers.push_back(std::move(wrapper));
+        return make_nb_class<PyArrayResultHandler>(
+            self.aval(), self.sharding(), self.committed(), self.skip_checks(),
+            std::move(wrappers));
+      });
 
   return absl::OkStatus();
 }

jaxlib/py_array.h

@@ -359,13 +359,22 @@ class PyArray : public nanobind::object {
 class PyArrayResultHandler {
  public:
   PyArrayResultHandler(nanobind::object aval, nanobind::object sharding,
-                       bool committed, bool skip_checks);
+                       bool committed, bool skip_checks,
+                       std::vector<nanobind::callable> wrappers = {});
 
-  PyArray Call(absl::Span<const PyArray> py_arrays) const;
-  PyArray Call(PyArray py_array) const;
+  nanobind::object Call(absl::Span<const PyArray> py_arrays) const;
+  nanobind::object Call(PyArray py_array) const;
 
-  PyArray Call(nb_class_ptr<PyClient> py_client, xla::ifrt::ArrayRef ifrt_array,
-               xla::Future<> result_status = xla::Future<>()) const;
+  nanobind::object Call(nb_class_ptr<PyClient> py_client,
+                        xla::ifrt::ArrayRef ifrt_array,
+                        xla::Future<> result_status = xla::Future<>()) const;
+
+  const std::vector<nanobind::callable>& wrappers() const { return wrappers_; }
+
+  nanobind::object aval() const { return aval_; }
+  nanobind::object sharding() const { return sharding_; }
+  bool committed() const { return committed_; }
+  bool skip_checks() const { return skip_checks_; }
 
  private:
   nanobind::object aval_;
@@ -376,6 +385,7 @@ class PyArrayResultHandler {
 
   xla::nb_dtype dtype_;
   std::vector<int64_t> shape_;
+  std::vector<nanobind::callable> wrappers_;
 };
 
 absl::StatusOr<nanobind::object> CudaArrayInterfaceToBuffer(

jaxlib/xla_client.py

@@ -47,7 +47,7 @@
 # Please suffix the version number with a brief description of your change
 # in a comment. The goal here is to force a merge conflict if two changes
 # attempt to grab the same version number.
-_version = 389  # LoadedExecutable.serialize
+_version = 390  # ResultHandler.wrap
 
 # An internal increasing version number for protecting jaxlib code against
 # ifrt changes.

tests/dtypes_test.py

@@ -34,6 +34,7 @@
 from jax._src import literals
 from jax._src import test_util as jtu
 from jax._src.lax import lax as lax_internal
+from jax._src.lib import jaxlib_extension_version
 
 config.parse_flags_with_absl()
 
@@ -872,6 +873,8 @@ def global_sharded_result_handler(aval, out_sharding, committed):
         phys_aval = core.physical_aval(aval)
         phys_handler_maker = pxla.global_result_handlers[core.ShapedArray]
         phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed)
+        if jaxlib_extension_version >= 390:
+          return phys_handler.wrap(lambda arr: earray.EArray(aval, arr))
         return lambda bufs: earray.EArray(aval, phys_handler(bufs))
 
     @dataclasses.dataclass(frozen=True)

tests/lax_test.py

@@ -48,6 +48,7 @@
 from jax._src.internal_test_util import lax_test_util
 from jax._src.lax import lax as lax_internal
 from jax._src.lax import utils as lax_utils
+from jax._src.lib import jaxlib_extension_version
 from jax._src.util import safe_zip
 from jax._src.tree_util import tree_map
 
@@ -3992,14 +3993,13 @@ def handler(_, buf):
 
   @staticmethod
   def global_sharded_result_handler(aval, out_sharding, committed):
-    def handler(arr):
-      from jax._src.array import ArrayImpl
-      if isinstance(arr, ArrayImpl):
-        buf, = arr._arrays
-      else:
-        buf, = arr
-      return FooArray(aval.shape, buf)
-    return handler
+    phys_sharding = out_sharding  # unlike KeyTyRules, assume same shape
+    phys_aval = core.physical_aval(aval)
+    phys_handler_maker = pxla.global_result_handlers[core.ShapedArray]
+    phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed)
+    if jaxlib_extension_version >= 390:
+      return phys_handler.wrap(lambda arr: FooArray(aval.shape, arr))
+    return lambda bufs: FooArray(aval.shape, phys_handler(bufs))
 
 
 class FooTy(dtypes.ExtendedDType):