Merge branch 'main' of https://github.com/google/jax

Author: nitins17

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

jax/_src/core.py

@@ -955,6 +955,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 +972,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 +1050,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
 

jax/_src/extend/ffi.py

@@ -18,7 +18,7 @@
 import ctypes
 import functools
 import os
-from typing import Any
+from typing import Any, overload
 
 import numpy as np
 
@@ -240,6 +240,43 @@ def _convert_layouts_for_ffi_call(
       for aval, layout in zip(avals, layouts))
 
 
+# ffi_call() returns as many results as result_shape_dtypes.
+@overload
+def ffi_call(
+    target_name: str,
+    result_shape_dtypes: ResultMetadata,
+    *deprecated_args: ArrayLike,
+    has_side_effect: bool = ...,
+    vmap_method: str | None = ...,
+    input_layouts: Sequence[FfiLayoutOptions] | None = ...,
+    output_layouts: FfiLayoutOptions | Sequence[FfiLayoutOptions] | None = ...,
+    input_output_aliases: dict[int, int] | None = ...,
+    custom_call_api_version: int = ...,
+    legacy_backend_config: str | None = ...,
+    vectorized: bool | DeprecatedArg = ...,
+    **deprecated_kwargs: Any,
+) -> Callable[..., Array] | Array:
+  ...
+
+
+@overload
+def ffi_call(
+    target_name: str,
+    result_shape_dtypes: Sequence[ResultMetadata],
+    *deprecated_args: ArrayLike,
+    has_side_effect: bool = ...,
+    vmap_method: str | None = ...,
+    input_layouts: Sequence[FfiLayoutOptions] | None = ...,
+    output_layouts: FfiLayoutOptions | Sequence[FfiLayoutOptions] | None = ...,
+    input_output_aliases: dict[int, int] | None = ...,
+    custom_call_api_version: int = ...,
+    legacy_backend_config: str | None = ...,
+    vectorized: bool | DeprecatedArg = ...,
+    **deprecated_kwargs: Any,
+) -> Callable[..., Sequence[Array]] | Sequence[Array]:
+  ...
+
+
 def ffi_call(
     target_name: str,
     result_shape_dtypes: ResultMetadata | Sequence[ResultMetadata],

jax/_src/interpreters/batching.py

@@ -596,9 +596,10 @@ def _batch_inner(axis_data, out_dim_dests, tag, in_dims, *in_vals):
     idx = memoize(lambda: BatchTracer(trace, make_iota(axis_data.size), 0,
                                       source_info_util.current()))
     in_tracers = map(partial(to_elt, trace, idx), in_vals, in_dims)
-    with core.set_current_trace(trace):
-      with core.extend_axis_env_nd([(axis_data.name, axis_data.size)]):
-        outs = yield in_tracers, {}
+    with (core.set_current_trace(trace),
+          core.extend_axis_env_nd([(axis_data.name, axis_data.size)]),
+          core.add_spmd_axis_names(axis_data.spmd_name)):
+      outs = yield in_tracers, {}
 
   out_dim_dests = out_dim_dests() if callable(out_dim_dests) else out_dim_dests
   out_vals = map(partial(from_elt, trace, axis_data.size), range(len(outs)),
@@ -795,9 +796,10 @@ def _batch_jaxpr_inner(axis_data, tag, in_axes, *in_vals):
     _, in_axes = resolve_ragged_axes(in_vals, in_axes)
     in_tracers = [BatchTracer(trace, val, dim) if dim is not None else val
                   for val, dim in zip(in_vals, in_axes)]
-    with core.set_current_trace(trace):
-      with core.extend_axis_env_nd([(axis_data.name, axis_data.size)]):
-        outs = yield in_tracers, {}
+    with (core.set_current_trace(trace),
+          core.extend_axis_env_nd([(axis_data.name, axis_data.size)]),
+          core.add_spmd_axis_names(axis_data.spmd_name)):
+      outs = yield in_tracers, {}
     out_vals, out_axes = unzip2(map(trace.to_batch_info, outs))
     new_out_axes = indirectify_ragged_axes_against_inputs_outputs(
         out_axes, in_vals, out_vals)

jax/_src/prng.py

@@ -279,6 +279,11 @@ def copy(self):
   __hash__ = None  # type: ignore[assignment]
   __array_priority__ = 100
 
+  def __array__(self, dtype: np.dtype | None = None, copy: bool | None = None) -> np.ndarray:
+    raise TypeError("JAX array with PRNGKey dtype cannot be converted to a NumPy array."
+                    " Use jax.random.key_data(arr) if you wish to extract the underlying"
+                    " integer array.")
+
   # Overwritten immediately below
   @property
   def at(self)                  -> _IndexUpdateHelper: assert False  # type: ignore[override]

jax/experimental/shard_map.py

@@ -1506,7 +1506,7 @@ def _unmentioned2(mesh: Mesh, names: AxisNames) -> list[AxisName]:
   # We use a filtered-down version of unmentioned to avoid defensive-psum over
   # more chips than required in the transpose-no-check-rep case.
   name_set = {n for ns in names.values() for n in ns}
-  return [n for n in mesh.axis_names if n not in name_set]
+  return [n for n in _all_mesh_names_except_spmd(mesh) if n not in name_set]
 
 
 def _shard_map_transpose(out_cts, *args, jaxpr, mesh, in_names, out_names,
@@ -1652,10 +1652,7 @@ def _pe_custom_params(unks_in, inst_in, kept_outs_known, kept_outs_staged,
 
 # TODO(mattjj): remove this mechanism when we revise mesh scopes
 def _all_mesh_names_except_spmd(mesh: Mesh, trace=None) -> tuple[AxisName, ...]:
-  trace = core.unsafe_get_current_trace() if trace is None else trace
-  stack = core.unsafe_get_trace_stack(trace)
-  batch_traces = [t for t in stack if isinstance(t, batching.BatchTrace)]
-  spmd_names = {n for trace in batch_traces for n in trace.axis_data.spmd_name }
+  spmd_names = core.get_axis_env().spmd_axis_names
   return tuple(name for name in mesh.axis_names if name not in spmd_names)
 
 # DCE

jaxlib/mosaic/dialect/gpu/BUILD

@@ -29,7 +29,9 @@ td_library(
     srcs = ["mosaic_gpu.td"],
     includes = ["."],
     deps = [
+        "@llvm-project//mlir:BasicPtxBuilderIntTdFiles",
         "@llvm-project//mlir:BuiltinDialectTdFiles",
+        "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
         "@llvm-project//mlir:OpBaseTdFiles",
     ],
 )
@@ -109,16 +111,19 @@ cc_library(
     hdrs = ["mosaic_gpu.h"],
     deps = [
         ":mosaic_gpu_inc_gen",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
         "@llvm-project//mlir:LLVMCommonConversion",
         "@llvm-project//mlir:LLVMDialect",
         "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:MemRefUtils",
         "@llvm-project//mlir:SCFUtils",
         "@llvm-project//mlir:Support",
         "@tsl//tsl/platform:statusor",
@@ -152,12 +157,19 @@ cc_test(
 gentbl_filegroup(
     name = "mosaic_gpu_python_gen_raw",
     tbl_outs = [
+        (
+            [
+                "-gen-python-enum-bindings",
+                "-bind-dialect=mosaic_gpu",
+            ],
+            "_mosaic_gpu_gen_enums_raw.py",
+        ),
         (
             [
                 "-gen-python-op-bindings",
                 "-bind-dialect=mosaic_gpu",
             ],
-            "_mosaic_gpu_gen_raw.py",
+            "_mosaic_gpu_gen_ops_raw.py",
         ),
     ],
     tblgen = "@llvm-project//mlir:mlir-tblgen",
@@ -169,10 +181,19 @@ gentbl_filegroup(
 )
 
 genrule(
-    name = "mosaic_gpu_python_gen",
-    srcs = ["_mosaic_gpu_gen_raw.py"],
-    outs = ["_mosaic_gpu_gen.py"],
-    cmd = "cat $(location _mosaic_gpu_gen_raw.py) | sed -e 's/^from \\./from jaxlib\\.mlir\\.dialects\\./g' > $@",
+    name = "mosaic_gpu_python_gen_enums",
+    srcs = ["_mosaic_gpu_gen_enums_raw.py"],
+    outs = ["_mosaic_gpu_gen_enums.py"],
+    cmd = """
+      cat $(location _mosaic_gpu_gen_enums_raw.py) | \
+      sed -e 's/^from \\.\\.ir/from jaxlib\\.mlir\\.ir/g; s/^from \\./from jaxlib\\.mlir\\.dialects\\./g' > $@""",
+)
+
+genrule(
+    name = "mosaic_gpu_python_gen_ops",
+    srcs = ["_mosaic_gpu_gen_ops_raw.py"],
+    outs = ["_mosaic_gpu_gen_ops.py"],
+    cmd = "cat $(location _mosaic_gpu_gen_ops_raw.py) | sed -e 's/^from \\./from jaxlib\\.mlir\\.dialects\\./g' > $@",
 )
 
 DIALECT_CAPI_SOURCES = [

jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc

@@ -18,18 +18,17 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
-#include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/string_view.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"  // IWYU pragma: keep
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/LogicalResult.h"
 #include "mlir/Conversion/LLVMCommon/MemRefBuilder.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -44,6 +43,12 @@ limitations under the License.
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Support/LLVM.h"
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "mlir/include/mlir/IR/Diagnostics.h"
 #include "tsl/platform/statusor.h"
 
 // Generated definitions.
@@ -232,11 +237,89 @@ void DeclareRuntimeFunctions(mlir::OpBuilder& builder) {
       .setVisibility(mlir::func::FuncOp::Visibility::Private);
 }
 
+bool IsContiguous(mlir::MemRefType type) {
+  return type.getLayout().isIdentity() ||
+         (type.hasStaticShape() && type.getNumElements() > 0 &&
+          mlir::memref::isStaticShapeAndContiguousRowMajor(type));
+}
+
+namespace {
+llvm::LogicalResult VerifyCommonLoadStoreOp(
+    mlir::Location loc, mlir::MemRefType gmem_type, absl::string_view gmem_name,
+    mlir::MemRefType smem_type, absl::string_view smem_name,
+    mlir::ArrayRef<int64_t> slice_lengths, int num_indices) {
+  auto error = [loc](auto... params) {
+    return emitError(loc, llvm::formatv(params...));
+  };
+
+  if (!IsContiguous(smem_type)) {
+    return error("The `{0}` memref must be contiguous.", smem_name);
+  }
+  if (gmem_type.getElementType() != smem_type.getElementType()) {
+    return error(
+        "The `source` and `destination` memrefs must have the same element "
+        "type.");
+  }
+  if (absl::c_any_of(slice_lengths, [](int64_t s) { return s < -1; })) {
+    return error(
+        "The `slice_lengths` attribute must not contain values less than -1.");
+  }
+  if (gmem_type.getRank() !=
+      smem_type.getRank() + absl::c_count(slice_lengths, -1)) {
+    return error(
+        "The rank of the `{0}` must be equal to the rank of the "
+        "`{1}` plus the number of collapsed dimensions as indicated "
+        "by -1 values in `slice_lengths`.",
+        gmem_name, smem_name);
+  }
+  if (num_indices != gmem_type.getRank()) {
+    return error("The size of `indices` must be equal to the rank of `{0}`.",
+                 gmem_name);
+  }
+  if (slice_lengths.size() != gmem_type.getRank()) {
+    return error(
+        "The size of `slice_lengths` must be equal to the rank of `{0}`.",
+        gmem_name);
+  }
+  return llvm::success();
+}
+}  // namespace
+
+llvm::LogicalResult AsyncLoadOp::verify() {
+  auto r = VerifyCommonLoadStoreOp(getLoc(), getSource().getType(), "source",
+                                   getDestination().getType(), "destination",
+                                   getSliceLengths(), getIndices().size());
+  if (failed(r)) {
+    return r;
+  }
+
+  for (int i = 0; i < getCollective().size(); ++i) {
+    for (int k = i + 1; k < getCollective().size(); ++k)
+      if (getCollective()[i] == getCollective()[k]) {
+        return emitError(
+            "The `collective` attribute must not contain duplicate "
+            "dimensions.");
+      }
+  }
+
+  return llvm::success();
+}
+
+llvm::LogicalResult AsyncStoreOp::verify() {
+  return VerifyCommonLoadStoreOp(getLoc(), getDestination().getType(),
+                                 "destination", getSource().getType(), "source",
+                                 getSliceLengths(), getIndices().size());
+}
+
 void MosaicGPUDialect::initialize() {
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "jaxlib/mosaic/dialect/gpu/mosaic_gpu_types.cc.inc"
       >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "jaxlib/mosaic/dialect/gpu/mosaic_gpu_attrdefs.cc.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "jaxlib/mosaic/dialect/gpu/mosaic_gpu_ops.cc.inc"