Merge pull request #131 from ROCm/ci-upstream-sync-11_1

CI: 11/06/24 upstream sync

Author: charleshofer

.github/workflows/ci-build.yaml

@@ -34,7 +34,12 @@ jobs:
         uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
         with:
           python-version: 3.11
-      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd  # v3.0.1
+      - run: python -m pip install pre-commit
+      - uses: actions/cache@3624ceb22c1c5a301c8db4169662070a689d9ea8 # v4.1.1
+        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
 
   build:
     name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ubuntu-20.04, x64=${{ matrix.enable-x64}})"
@@ -211,14 +216,14 @@ jobs:
 
   ffi:
     name: FFI example
-    runs-on: ubuntu-latest
-    timeout-minutes: 5
+    runs-on: ROCM-Ubuntu
+    timeout-minutes: 30
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-    - name: Set up Python 3.11
+    - name: Set up Python
       uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
       with:
-        python-version: 3.11
+        python-version: 3.12
     - name: Get pip cache dir
       id: pip-cache
       run: |
@@ -239,6 +244,10 @@ jobs:
         # a different toolchain. GCC is the default compiler on the
         # 'ubuntu-latest' runner, but we still set this explicitly just to be
         # clear.
-        CMAKE_ARGS: -DCMAKE_C_COMPILER=gcc -DCMAKE_CXX_COMPILER=g++
-    - name: Run tests
+        CMAKE_ARGS: -DCMAKE_CXX_COMPILER=g++ #-DJAX_FFI_EXAMPLE_ENABLE_CUDA=ON
+    - name: Run CPU tests
+      run: python -m pytest examples/ffi/tests
+      env:
+        JAX_PLATFORM_NAME: cpu
+    - name: Run GPU tests
       run: python -m pytest examples/ffi/tests

.github/workflows/jax-array-api.yml

@@ -28,7 +28,7 @@ jobs:
       with:
         repository: data-apis/array-api-tests
         # TODO(jakevdp) update this to a stable release/tag when available.
-        ref: 'b4c0823469c02d6ce6e512ad4c2bd8ba42b1b4b2'  # Latest commit as of 2024-09-09
+        ref: 'bcd5919bbbdf4d4806b5b2613b4d8c0bc0625c54'  # Latest commit as of 2024-10-31 👻
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python ${{ matrix.python-version }}

CHANGELOG.md

@@ -23,6 +23,26 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * The deprecated module `jax.experimental.export` has been removed. It was replaced
     by {mod}`jax.export` in JAX v0.4.30. See the [migration guide](https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export)
     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.
+  * The following deprecated methods and functions in {mod}`jax.export` have
+    been removed:
+      * `jax.export.DisabledSafetyCheck.shape_assertions`: it had no effect
+        already.
+      * `jax.export.Exported.lowering_platforms`: use `platforms`.
+      * `jax.export.Exported.mlir_module_serialization_version`:
+        use `calling_convention_version`.
+      * `jax.export.Exported.uses_shape_polymorphism`:
+         use `uses_global_constants`.
+      * the `lowering_platforms` kwarg for {func}`jax.export.export`: use
+        `platforms` instead.
+  * Hashing of tracers, which has been deprecated since version 0.4.30, now
+    results in a `TypeError`.
+
+* New Features
+  * {func}`jax.jit` got a new `compiler_options: dict[str, Any]` argument, for
+    passing compilation options to XLA. For the moment it's undocumented and
+    may be in flux.
 
 ## jax 0.4.35 (Oct 22, 2024)
 

README.md

@@ -411,23 +411,18 @@ community-supported conda build, and answers to some frequently-asked questions.
 
 ## Neural network libraries
 
-Multiple Google research groups develop and share libraries for training neural
-networks in JAX. If you want a fully featured library for neural network
+Multiple Google research groups at Google DeepMind and Alphabet develop and share libraries
+for training neural networks in JAX. If you want a fully featured library for neural network
 training with examples and how-to guides, try
-[Flax](https://github.com/google/flax). Check out the new [NNX](https://flax.readthedocs.io/en/latest/nnx/index.html) API for a
-simplified development experience.
-
-Google X maintains the neural network library
-[Equinox](https://github.com/patrick-kidger/equinox). This is used as the
-foundation for several other libraries in the JAX ecosystem.
-
-In addition, DeepMind has open-sourced an [ecosystem of libraries around
-JAX](https://deepmind.com/blog/article/using-jax-to-accelerate-our-research)
-including [Optax](https://github.com/deepmind/optax) for gradient processing and
-optimization, [RLax](https://github.com/deepmind/rlax) for RL algorithms, and
-[chex](https://github.com/deepmind/chex) for reliable code and testing. (Watch
-the NeurIPS 2020 JAX Ecosystem at DeepMind talk
-[here](https://www.youtube.com/watch?v=iDxJxIyzSiM))
+[Flax](https://github.com/google/flax) and its [documentation site](https://flax.readthedocs.io/en/latest/nnx/index.html).
+
+Check out the [JAX Ecosystem section](https://jax.readthedocs.io/en/latest/#ecosystem)
+on the JAX documentation site for a list of JAX-based network libraries, which includes
+[Optax](https://github.com/deepmind/optax) for gradient processing and
+optimization, [chex](https://github.com/deepmind/chex) for reliable code and testing, and
+[Equinox](https://github.com/patrick-kidger/equinox) for neural networks.
+(Watch the NeurIPS 2020 JAX Ecosystem at DeepMind talk
+[here](https://www.youtube.com/watch?v=iDxJxIyzSiM) for additional details.)
 
 ## Citing JAX
 

build/requirements.in

@@ -3,11 +3,6 @@
 #
 -r test-requirements.txt
 
-# matplotlib 3.9.0 pins NumPy 1.23, which is incompatible with the requirement
-# below.
-matplotlib~=3.8.4; python_version<="3.10"
-matplotlib; python_version>="3.11"
-
 #
 # build deps
 #