Merge pull request #185 from ROCm/ci-upstream-sync-12-12-2024

CI: 12/12/24 upstream sync

Author: charleshofer

.bazelrc

@@ -96,6 +96,11 @@ build:avx_windows --copt=/arch:AVX
 
 build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1
 
+# Config setting to build oneDNN with Compute Library for the Arm Architecture (ACL).
+build:mkl_aarch64_threadpool --define=build_with_mkl_aarch64=true
+build:mkl_aarch64_threadpool --@compute_library//:openmp=false
+build:mkl_aarch64_threadpool -c opt
+
 # Disable clang extention that rejects type definitions within offsetof.
 # This was added in clang-16 by https://reviews.llvm.org/D133574.
 # Can be removed once upb is updated, since a type definition is used within

.github/workflows/bazel_cpu_rbe.yml

@@ -11,6 +11,9 @@ on:
         options:
         - 'yes'
         - 'no'
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -21,15 +24,18 @@ jobs:
     if: github.event.repository.fork == false
     strategy:
       matrix:
-        runner: ["linux-x86-n2-16", "linux-arm64-t2a-16"]
+        runner: ["linux-x86-n2-16", "linux-arm64-c4a-16"]
+        enable-x_64: [1, 0]
 
     runs-on: ${{ matrix.runner }}
-    # TODO(b/369382309): Replace Linux Arm64 container with the ml-build container once it is available
     container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
-                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') }}
+                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') }}
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
+      JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
+
+    name: "Bazel CPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
 
     steps:
       - uses: actions/checkout@v3

.github/workflows/bazel_gpu_rbe.yml

@@ -11,6 +11,9 @@ on:
         options:
         - 'yes'
         - 'no'
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -22,12 +25,16 @@ jobs:
     strategy:
       matrix:
         runner: ["linux-x86-n2-16"]
+        enable-x_64: [1, 0]
 
     runs-on: ${{ matrix.runner }}
     container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
+      JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
+
+    name: "Bazel single accelerator GPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
 
     steps:
       - uses: actions/checkout@v3

.pre-commit-config.yaml

@@ -36,7 +36,7 @@ repos:
   - id: mypy
     files: (jax/|tests/typing_test\.py)
     exclude: jax/_src/basearray.py|jax/numpy/__init__.py  # Use pyi instead
-    additional_dependencies: [types-requests==2.31.0, jaxlib, numpy~=2.1.0]
+    additional_dependencies: [types-requests==2.31.0, jaxlib, numpy>=2.2.0]
     args: [--config=pyproject.toml]
 
 - repo: https://github.com/mwouts/jupytext

CHANGELOG.md

@@ -19,6 +19,11 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     APIs of the same name in {mod}`jax.extend.core`; see the documentation for
     {mod}`jax.extend` for information on the compatibility guarantees of these
     semi-public extensions.
+  * Several previously-deprecated APIs have been removed, including:
+    * from {mod}`jax.core`: `check_eqn`, `check_type`,  `check_valid_jaxtype`, and
+      `non_negative_dim`.
+    * from {mod}`jax.lib.xla_bridge`: `xla_client` and `default_backend`.
+    * from {mod}`jax.lib.xla_client`: `_xla` and `bfloat16`.
 
 ## jax 0.4.37 (Dec 9, 2024)
 

README.md

@@ -47,7 +47,7 @@ are instances of such transformations. Others are
 [`pmap`](#spmd-programming-with-pmap) for single-program multiple-data (SPMD)
 parallel programming of multiple accelerators, with more to come.
 
-This is a research project, not an official Google product. Expect bugs and
+This is a research project, not an official Google product. Expect
 [sharp edges](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html).
 Please help by trying it out, [reporting
 bugs](https://github.com/jax-ml/jax/issues), and letting us know what you

build/build.py

@@ -485,7 +485,10 @@ async def main():
 
     if not args.disable_mkl_dnn:
       logging.debug("Enabling MKL DNN")
-      wheel_build_command.append("--config=mkl_open_source_only")
+      if target_cpu == "aarch64":
+          wheel_build_command.append("--config=mkl_aarch64_threadpool")
+      else:
+          wheel_build_command.append("--config=mkl_open_source_only")
 
     if args.target_cpu_features == "release":
       if arch in ["x86_64", "AMD64"]:

docs/jax.lib.rst

@@ -11,7 +11,6 @@ jax.lib.xla_bridge
 .. autosummary::
   :toctree: _autosummary
 
-  default_backend
   get_backend
   get_compile_options
 

docs/jax.numpy.rst

@@ -274,6 +274,7 @@ namespace; they are listed below.
     mask_indices
     matmul
     matrix_transpose
+    matvec
     max
     maximum
     mean
@@ -428,6 +429,7 @@ namespace; they are listed below.
     var
     vdot
     vecdot
+    vecmat
     vectorize
     vsplit
     vstack

examples/ffi/tests/cpu_examples_test.py

@@ -37,10 +37,10 @@ def test_array_attr_jit_cache(self):
     jit_array_attr = jax.jit(cpu_examples.array_attr, static_argnums=(0,))
     with jtu.count_jit_and_pmap_lowerings() as count:
       jit_array_attr(5)
-    self.assertEqual(count[0], 1)  # compiles once the first time
+    self.assertEqual(count(), 1)  # compiles once the first time
     with jtu.count_jit_and_pmap_lowerings() as count:
       jit_array_attr(5)
-    self.assertEqual(count[0], 0)  # cache hit
+    self.assertEqual(count(), 0)  # cache hit
 
   def test_array_attr_no_jit(self):
     with jax.disable_jit():