Squashed commit of the following:

commit 093c6e9 Merge: e1a77ee d0cb318 Author: Keshav <keshavb@nvidia.com> Date: Wed Sep 18 13:37:44 2024 -0700 Merge remote-tracking branch 'upstream/main' into disable_remat_pass commit e1a77ee Author: Keshav <keshavb@nvidia.com> Date: Wed Sep 18 13:35:37 2024 -0700 minor changes commit d0cb318 Merge: b51c653 bef36c4 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 13:34:11 2024 -0700 Merge pull request jax-ml#23736 from hawkinsp:changelog PiperOrigin-RevId: 676111400 commit b51c653 Merge: dbc03cf 57a4b76 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 13:33:05 2024 -0700 Merge pull request jax-ml#23737 from jakevdp:digitize-doc PiperOrigin-RevId: 676111220 commit dbc03cf Author: Dan Foreman-Mackey <danfm@google.com> Date: Wed Sep 18 12:39:58 2024 -0700 Re-land jax-ml#23261 with appropriate compatibility checks. PiperOrigin-RevId: 676092618 commit b164d67 Merge: cd04d0f 541b3a3 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 12:05:03 2024 -0700 Merge pull request jax-ml#23247 from kaixih:sliding_window_attn PiperOrigin-RevId: 676079831 commit 57a4b76 Author: Jake VanderPlas <jakevdp@google.com> Date: Wed Sep 18 11:59:00 2024 -0700 Improve documentation for jnp.digitize commit bef36c4 Author: Peter Hawkins <phawkins@google.com> Date: Wed Sep 18 18:57:03 2024 +0000 Add Python 3.13 wheels to changelog. commit cd04d0f Merge: 016c499 c756d9b Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 10:00:03 2024 -0700 Merge pull request jax-ml#23726 from hawkinsp:debug PiperOrigin-RevId: 676030839 commit 016c499 Author: Sergei Lebedev <slebedev@google.com> Date: Wed Sep 18 09:56:44 2024 -0700 Removed leftover usages of GPUGridSpec from Pallas Mosaic GPU tests PiperOrigin-RevId: 676029854 commit 9dd363d Author: Luke Baumann <lukebaumann@google.com> Date: Wed Sep 18 09:28:25 2024 -0700 Export `jax.lib.xla_extension.ifrt_programs`. PiperOrigin-RevId: 676020419 commit e27f1e9 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 09:03:55 2024 -0700 Change Python version 3.13.0rc2 to 3.13.0-rc.2. The value is taken from [the versions manifest](https://raw.githubusercontent.com/actions/python-versions/main/versions-manifest.json). PiperOrigin-RevId: 676012255 commit 442e863 Author: Sergei Lebedev <slebedev@google.com> Date: Wed Sep 18 08:56:49 2024 -0700 Added a missing branch to `mgpu.FragmentedArray.astype` Previously, an unsupported cast produced a `NameError` instead. PiperOrigin-RevId: 676010161 commit 6236b8f Merge: 826843a 1cc9661 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 08:57:38 2024 -0700 Merge pull request jax-ml#23667 from dfm:always-lower-jnp-dot-to-dot-general PiperOrigin-RevId: 676010154 commit c756d9b Author: Peter Hawkins <phawkins@google.com> Date: Wed Sep 18 15:44:45 2024 +0000 Fix error in debugger tests that is showing up in CI. I'm unsure why this started happening now, but sometimes we get an invalid offset for a frame. Be tolerant of that case. commit 826843a Merge: c191bbc 922e652 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 08:42:39 2024 -0700 Merge pull request jax-ml#23723 from hawkinsp:setuptools PiperOrigin-RevId: 676005613 commit c191bbc Author: Yash Katariya <yashkatariya@google.com> Date: Wed Sep 18 08:40:30 2024 -0700 Make `debug.print` work with static args. Fixes: jax-ml#23600 PiperOrigin-RevId: 676005582 commit 1cc9661 Author: Dan Foreman-Mackey <danfm@google.com> Date: Mon Sep 16 14:18:29 2024 -0400 Unconditionally lower jnp.dot to lax.dot_general. jax-ml#16721 added a condition to lower calls to `jnp.dot` with scalar inputs to `lax.mul` instead of `lax.dot_general`. AFAICT, jax-ml#16826 fixed the issue that this was solving, so this condition should no longer be necessary. Removing this condition simplifies the addition of new arguments to `dot` and `dot_general`, including the `algorithm` parameter that I am currently working on in jax-ml#23574, so now seemed like a good time to remove it! commit 922e652 Author: Peter Hawkins <phawkins@google.com> Date: Wed Sep 18 15:17:49 2024 +0000 Replace plat-name with plat_name. The former seems to elicit a deprecation warning from setuptools recently. commit 69ba060 Author: Dan Foreman-Mackey <danfm@google.com> Date: Wed Sep 18 07:40:58 2024 -0700 Reverts e15ec1e PiperOrigin-RevId: 675987338 commit 44a7f04 Merge: 0a29696 2834c13 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 07:31:00 2024 -0700 Merge pull request jax-ml#23708 from jakevdp:sort-complex PiperOrigin-RevId: 675983957 commit 0a29696 Merge: e15ec1e 73c38cb Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 07:08:24 2024 -0700 Merge pull request jax-ml#23698 from dfm:dev-clang-warning PiperOrigin-RevId: 675977448 commit 2834c13 Author: Jake VanderPlas <jakevdp@google.com> Date: Tue Sep 17 15:32:25 2024 -0700 jnp.sort_complex: fix output for N-dimensional inputs commit 73c38cb Author: Dan Foreman-Mackey <danfm@google.com> Date: Tue Sep 17 14:00:21 2024 -0400 Add a note to the developer docs making it clear that clang is the only toolchain that is actively supported for source compilation. As discussed in jax-ml#23687 commit e15ec1e Merge: 48d8fce 3f2bc9b Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 06:56:28 2024 -0700 Merge pull request jax-ml#23261 from joaospinto:stablehlo.tan PiperOrigin-RevId: 675973798 commit 48d8fce Merge: 4e6f690 2714469 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 06:54:28 2024 -0700 Merge pull request jax-ml#23563 from rajasekharporeddy:testbranch1 PiperOrigin-RevId: 675973225 commit 4e6f690 Merge: b7c91e9 611ad63 Author: jax authors <google-ml-automation@google.com> Date: Wed Sep 18 06:35:15 2024 -0700 Merge pull request jax-ml#23653 from apaszke:torchsaic PiperOrigin-RevId: 675967844 commit b7c91e9 Author: Sergei Lebedev <slebedev@google.com> Date: Wed Sep 18 06:22:14 2024 -0700 Lookup `shape` and `dtype` directly on `state.AbstractRef` instead of going through `inner_aval` This is just a cleanup. No behavior changes are expected. PiperOrigin-RevId: 675964703 commit 611ad63 Author: Adam Paszke <adam.paszke@gmail.com> Date: Fri Sep 6 16:09:58 2024 +0000 Add basic PyTorch integration for Mosaic GPU We have already had most of the relevant pieces and we only needed to connect them together. The most sensitive change is perhaps that I needed to expose one more symbol from the XLA GPU plugin, but I don't think it should be a problem. commit e903369 Author: Sergei Lebedev <slebedev@google.com> Date: Wed Sep 18 05:25:37 2024 -0700 Pulled `scratch_shapes` into `GridSpec` It is supported by Mosaic TPU and Mosaic GPU and unsupported by Triton. PiperOrigin-RevId: 675950199 commit 2714469 Author: rajasekharporeddy <rajasekharp@google.com> Date: Wed Sep 18 17:06:28 2024 +0530 Deprecate passing NdArrays with ndim != 1 and non-arraylike inputs to jnp.trim_zeros commit b904599 Author: Sergei Lebedev <slebedev@google.com> Date: Wed Sep 18 04:23:25 2024 -0700 `pl.debug_print` no longer restricts values to be scalars This allows printing arrays on Triton and soon on Mosaic GPU. PiperOrigin-RevId: 675935666 commit 988ed2b Author: jax authors <google-ml-automation@google.com> Date: Tue Sep 17 21:09:26 2024 -0700 Add support for SMEM windows in Pallas custom pipeline. PiperOrigin-RevId: 675822640 commit f79d85b Merge: 1b74cfd cc28d63 Author: Keshav <keshavb@nvidia.com> Date: Tue Sep 17 18:58:33 2024 -0700 Merge remote-tracking branch 'upstream/main' into disable_remat_pass commit cc28d63 Merge: 8bcdb12 9d3762b Author: jax authors <google-ml-automation@google.com> Date: Tue Sep 17 17:36:36 2024 -0700 Merge pull request jax-ml#23682 from sharadmv:pallas-async-docs PiperOrigin-RevId: 675770723 commit 1b74cfd Author: Keshav <keshavb@nvidia.com> Date: Tue Sep 17 17:23:30 2024 -0700 disable remat hlo pass by default commit 8bcdb12 Author: jax authors <google-ml-automation@google.com> Date: Tue Sep 17 16:50:55 2024 -0700 Add CI jobs for python 3.13.0rc2. PiperOrigin-RevId: 675758096 commit 8b5b717 Author: Yash Katariya <yashkatariya@google.com> Date: Tue Sep 17 16:39:55 2024 -0700 Fix jaxpr equation context propagation in jaxpr equations when `inline=True`. PiperOrigin-RevId: 675754808 commit 86fe463 Author: Parker Schuh <parkers@google.com> Date: Tue Sep 17 16:10:41 2024 -0700 [Take 2] Generalize global jit cpp cache keys so we can add more keys than the current donate_argnums. This allows us to get more cache hits globally. For example: Before: jax.jit(f, out_shardings=s)(arr) jax.jit(f, out_shardings=s)(arr) # cpp cache miss After: jax.jit(f, out_shardings=s)(arr) jax.jit(f, out_shardings=s)(arr) # cpp cache hit Reverts b615266 PiperOrigin-RevId: 675746175 commit e92a599 Author: Christos Perivolaropoulos <cperivol@google.com> Date: Tue Sep 17 15:26:42 2024 -0700 [mosaic_gpu] Better error message for misaligned tma_transpose with dtype. PiperOrigin-RevId: 675731295 commit 7864648 Merge: 3f2c58b 83a7555 Author: jax authors <google-ml-automation@google.com> Date: Tue Sep 17 15:12:50 2024 -0700 Merge pull request jax-ml#23679 from selamw1:docstring_sort_complex PiperOrigin-RevId: 675726527 commit 83a7555 Author: selamw1 <selamw@google.com> Date: Mon Sep 16 16:47:52 2024 -0700 docstring_sort_complex_added input_array_modified commit 9d3762b Author: Sharad Vikram <sharad.vikram@gmail.com> Date: Mon Sep 16 19:18:22 2024 -0700 [Pallas] Add design note for async ops on TPU commit 3f2bc9b Author: Joao Sousa-Pinto <joaospinto@gmail.com> Date: Mon Aug 26 17:25:16 2024 -0700 Lower tan to StableHLO instead of CHLO. Fixes jax-ml#23259 commit 541b3a3 Author: kaixih <kaixih@nvidia.com> Date: Mon Aug 26 17:32:38 2024 +0000 New feature
keshavb96 · Sep 18, 2024 · 9171e66 · 9171e66
1 parent 3f2c58b
commit 9171e66
Show file tree

Hide file tree

Showing 54 changed files with 1,536 additions and 308 deletions.
diff --git a/.bazelrc b/.bazelrc
@@ -215,6 +215,8 @@ build:rbe_cpu_linux_py3.11 --config=rbe_cpu_linux_base
 build:rbe_cpu_linux_py3.11 --repo_env HERMETIC_PYTHON_VERSION="3.11"
 build:rbe_cpu_linux_py3.12 --config=rbe_cpu_linux_base
 build:rbe_cpu_linux_py3.12 --repo_env HERMETIC_PYTHON_VERSION="3.12"
+build:rbe_cpu_linux_py3.13 --config=rbe_cpu_linux_base
+build:rbe_cpu_linux_py3.13 --repo_env HERMETIC_PYTHON_VERSION="3.13"
 
 build:rbe_linux_cuda_base --config=rbe_linux
 build:rbe_linux_cuda_base --config=cuda
@@ -237,6 +239,8 @@ build:rbe_linux_cuda12.3_nvcc_py3.11 --config=rbe_linux_cuda12.3_nvcc_base
 build:rbe_linux_cuda12.3_nvcc_py3.11 --repo_env HERMETIC_PYTHON_VERSION="3.11"
 build:rbe_linux_cuda12.3_nvcc_py3.12 --config=rbe_linux_cuda12.3_nvcc_base
 build:rbe_linux_cuda12.3_nvcc_py3.12 --repo_env HERMETIC_PYTHON_VERSION="3.12"
+build:rbe_linux_cuda12.3_nvcc_py3.13 --config=rbe_linux_cuda12.3_nvcc_base
+build:rbe_linux_cuda12.3_nvcc_py3.13 --repo_env HERMETIC_PYTHON_VERSION="3.13"
 
 # These you may need to change for your own GCP project.
 build:tensorflow_testing_rbe --project_id=tensorflow-testing

diff --git a/.github/workflows/wheel_win_x64.yml b/.github/workflows/wheel_win_x64.yml
@@ -17,7 +17,7 @@ jobs:
       matrix:
         os: [windows-2019-32core]
         arch: [AMD64]
-        pyver: ['3.10', '3.11', '3.12']
+        pyver: ['3.10', '3.11', '3.12', '3.13.0-rc.2']
     name: ${{ matrix.os }} ${{ matrix.pyver }} jaxlib wheel build
     runs-on: ${{ matrix.os }}
 

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,15 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 
 ## jax 0.4.34
 
+* New Functionality
+  * This release includes wheels for Python 3.13. Free-threading mode is not yet
+    supported.
+
+* Deprecations
+  * In {func}`jax.numpy.trim_zeros`, non-arraylike arguments or arraylike
+    arguments with `ndim != 1` are now deprecated, and in the future will result
+    in an error.
+
 * Deletion:
   * `jax.xla_computation` is deleted. It's been 3 months since it's deprecation
     in 0.4.30 JAX release.

diff --git a/build/requirements.in b/build/requirements.in
@@ -23,3 +23,5 @@ ml_dtypes>=0.4.0
 opt_einsum
 zstandard
 etils[epath]
+# TODO(ybaturina): remove setuptools version
+setuptools<71.0.0
diff --git a/build/requirements_lock_3_13.txt b/build/requirements_lock_3_13.txt
@@ -732,7 +732,9 @@ zstandard==0.23.0 \
     # via -r build/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==75.1.0 \
-    --hash=sha256:35ab7fd3bcd95e6b7fd704e4a1539513edad446c097797f2985e0e4b960772f2 \
-    --hash=sha256:d59a21b17a275fb872a9c3dae73963160ae079f1049ed956880cd7c09b120538
-    # via -r build/test-requirements.txt
+setuptools==70.3.0 \
+    --hash=sha256:f171bab1dfbc86b132997f26a119f6056a57950d058587841a0082e8830f9dc5 \
+    --hash=sha256:fe384da74336c398e0d956d1cae0669bc02eed936cdb1d49b57de1990dc11ffc
+    # via
+    #   -r build/requirements.in
+    #   -r build/test-requirements.txt
diff --git a/docs/developer.md b/docs/developer.md
@@ -31,23 +31,33 @@ guidance on pip installation (e.g., for GPU and TPU support).
 
 ### Building `jaxlib` from source
 
+```{warning}
+While it should typically be possible to compile `jaxlib` from source using
+most modern compilers, the builds are only tested using clang. Pull requests
+are welcomed to improve support for different toolchains, but other compilers
+are not actively supported.
+```
+
 To build `jaxlib` from source, you must also install some prerequisites:
 
-- a C++ compiler (g++, clang, or MSVC)
+- A C++ compiler:
 
-  On Ubuntu or Debian you can install the necessary prerequisites with:
+  As mentioned in the box above, it is best to use a recent version of clang
+  (at the time of writing, the version we test is 18), but other compilers (e.g.
+  g++ or MSVC) may work.
 
-  ```
-  sudo apt install g++ python python3-dev
-  ```
+  On Ubuntu or Debian you can follow the instructions from the
+  [LLVM](https://apt.llvm.org/) documentation to install the latest stable
+  version of clang.
 
   If you are building on a Mac, make sure XCode and the XCode command line tools
   are installed.
 
   See below for Windows build instructions.
 
-- there is no need to install Python dependencies locally, as your system
-  Python will be ignored during the build; please check
+- Python: for running the build helper script. Note that there is no need to
+  install Python dependencies locally, as your system Python will be ignored
+  during the build; please check
   [Managing hermetic Python](#managing-hermetic-python) for details.
 
 To build `jaxlib` for CPU or TPU, you can run:
@@ -86,7 +96,7 @@ the `build/build.py` script itself will be processed by your system Python
 interpreter. By default, the wheel is written to the `dist/` subdirectory of the
 current directory.
 
-*  JAX versions starting from v.0.4.32: you can provide custom CUDA and CUDNN 
+*  JAX versions starting from v.0.4.32: you can provide custom CUDA and CUDNN
    versions in the configuration options. Bazel will download them and use as
    target dependencies.
 
@@ -259,8 +269,8 @@ together with their corresponding hashes are specified in
 `build/requirements_lock_<python version>.txt` files (
 e.g. `build/requirements_lock_3_12.txt` for `Python 3.12`).
 
-To update the lock files, make sure `build/requirements.in` contains the desired 
-direct dependencies list and then execute the following command (which will call 
+To update the lock files, make sure `build/requirements.in` contains the desired
+direct dependencies list and then execute the following command (which will call
 [pip-compile](https://pypi.org/project/pip-tools/) under the hood):
 
 ```
@@ -382,7 +392,7 @@ sudo apt-get install libopenblas-dev -y
    example for `Python 3.13` it should have something
    like `"3.13": "//build:requirements_lock_3_13.txt"`. Note, the key in the
    `requirements` parameter must always be in `"major.minor"` version format, so
-   even if you are building Python version `3.13.0rc1` the corresponding 
+   even if you are building Python version `3.13.0rc1` the corresponding
    `requirements` entry must still be `"3.13": "//build:requirements_lock_3_13.txt"`,
    **not** `"3.13.0rc1": "//build:requirements_lock_3_13_0rc1.txt"`.
 

diff --git a/docs/pallas/CHANGELOG.md b/docs/pallas/CHANGELOG.md
@@ -11,15 +11,31 @@ For the overall JAX change log see [here](https://jax.readthedocs.io/en/latest/c
 Remember to align the itemized text with the first line of an item within a list.
 -->
 
-## Released with jax 0.4.32
+## Released with jax 0.4.34
 
 * Changes
-  * The kernel function is not allowed to close over constants. Instead, all the needed arrays
-    must be passed as inputs, with proper block specs ({jax-issue}`#22746`).
+
+  * {func}`jax.experimental.pallas.debug_print` no longer requires all arguments
+    to be scalars. The restrictions on the arguments are backend-specific:
+    Non-scalar arguments are currently only supported on GPU, when using Triton.
 
 * Deprecations
 
-* New functionality:
+* New functionality
+
+  * {func}`jax.experimental.pallas.pallas_call` now accepts `scratch_shapes`,
+    a PyTree specifying backend-specific temporary objects needed by the
+    kernel, for example, buffers, synchronization primitives etc.
+
+## Released with jax 0.4.33 (September 16, 2024)
+
+## Released with jax 0.4.32 (September 11, 2024)
+
+* Changes
+  * The kernel function is not allowed to close over constants. Instead, all the needed arrays
+    must be passed as inputs, with proper block specs ({jax-issue}`#22746`).
+
+* New functionality
   * Improved error messages for mistakes in the signature of the index map functions,
     to include the name and source location of the index map.
 
@@ -44,10 +60,6 @@ Remember to align the itemized text with the first line of an item within a list
   * Previously it was possible to import many APIs that are meant to be
     private, as `jax.experimental.pallas.pallas`. This is not possible anymore.
 
-
-* Deprecations
-
-
 * New Functionality
   * Added documentation for BlockSpec: {ref}`pallas_grids_and_blockspecs`.
   * Improved error messages for the {func}`jax.experimental.pallas.pallas_call`
@@ -73,7 +85,3 @@ Remember to align the itemized text with the first line of an item within a list
   * Added checkify support for {func}`jax.experimental.pallas.pallas_call` in
     interpret mode ({jax-issue}`#21862`).
   * Improved support for PRNG keys for TPU kernels ({jax-issue}`#21773`).
-
-
-
-