These 16 commits are when the Protocol Buffers files have changed:
| Commit: | 6d68026 | |
|---|---|---|
| Author: | Olli Lupton | |
| Committer: | Olli Lupton | |
Mosaic GPU: pre-launch check for multicast support This gives a much better error message from kernels using multimem PTX instructions on incompatible systems.
| Commit: | 9fa9c3b | |
|---|---|---|
| Author: | Levon Ter-Grigoryan | |
| Committer: | jax authors | |
[Mosaic:GPU] Rename is_comm_used as is_nvshmem_used PiperOrigin-RevId: 878508055
| Commit: | a3ffef5 | |
|---|---|---|
| Author: | Aliia Khasanova | |
| Committer: | jax authors | |
Add serialization support for Mosaic GPU kernels. This change introduces a proto definition to serialize the state of a compiled Mosaic GPU kernel, including the object file and function names. The `CustomCallResources` struct is updated to stroe the kernel hash and include serialization and deserialization methods, allowing it to be used with XLA FFI's type registry. Also introduces `GetOrCreateKernel` to deduplicate kernels based on their kernel_hash (preserved during de/serialization). PiperOrigin-RevId: 877392204
| Commit: | 91223b9 | |
|---|---|---|
| Author: | Ilya Tikhonovskiy | |
| Committer: | jax authors | |
Follow up pr after triton integration cl and tokamax pr. Remove obsolete cluster dimensions from TritonKernel proto. PiperOrigin-RevId: 856837148
| Commit: | 81fe5dd | |
|---|---|---|
| Author: | Ilya Tikhonovskiy | |
| Committer: | jax authors | |
Integrate Triton up to 8d445186 https://github.com/openxla/triton/tree/triton_integrate_branch-1.15 PiperOrigin-RevId: 843206657
| Commit: | 72357ce | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
[JAX] Migrate most remaining C++ code in jax to the jax:: namespace. PiperOrigin-RevId: 790476584
The documentation is generated from this commit.
| Commit: | dd1eac8 | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
[JAX] Migrate most remaining C++ code in jax to the jax:: namespace. PiperOrigin-RevId: 789738942
The documentation is generated from this commit.
| Commit: | 39d7c38 | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
Move contents of jaxlib/xla into jaxlib/ Having the directory structure of the jaxlib wheel be different to the source tree confuses type checkers such as mypy, since sometimes they find type stubs in the installed jaxlib wheel, and sometimes from the installed source tree. Instead: * don't include type stubs in the jaxlib wheel * don't install the jaxlib wheel as part of pre-commit * make sure that the location of type stubs (and the underlying libraries) is in the same position in the `jaxlib/` directory of the JAX source tree as it would be for the jaxlib wheel when installed. For now, we leave some stubs that forward from the old locations to the new locations for certain headers and modules. These will be removed after migrating some users. PiperOrigin-RevId: 750650528
| Commit: | f02a204 | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
Move contents of jaxlib/xla into jaxlib/ Having the directory structure of the jaxlib wheel be different to the source tree is confusing type checkers such as mypy. PiperOrigin-RevId: 750398365
| Commit: | 588b693 | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
[JAX] [XLA:Python] Migrate more Python modules to JAX. PiperOrigin-RevId: 740507886
| Commit: | 51560bf | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
[JAX] [XLA:Python] Migrate pytree module to JAX. PiperOrigin-RevId: 740142231
| Commit: | ab3c1b5 | |
|---|---|---|
| Author: | jax authors | |
| Committer: | jax authors | |
[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1 PiperOrigin-RevId: 599809560
| Commit: | 3bb41c3 | |
|---|---|---|
| Author: | Peter Hawkins | |
| Committer: | jax authors | |
Move jaxlib/ under jax/_src/lib. Prototype only, not ready for review. PiperOrigin-RevId: 570062727
| Commit: | 4ac2bdc | |
|---|---|---|
| Author: | Chris Jones | |
| Committer: | jax authors | |
[jax_triton] Add user-specified `name` field to serialized format. PiperOrigin-RevId: 557415723
| Commit: | 714156d | |
|---|---|---|
| Author: | Chris Jones | |
| Committer: | jax authors | |
[jax_triton] Add support for float scalar inputs. Python `float`s are inferred as "f64". Values can be passed as "f32" using `np.float32(value)`. PiperOrigin-RevId: 552036612
| Commit: | f238667 | |
|---|---|---|
| Author: | Chris Jones | |
| Committer: | jax authors | |
Make JAX-Triton calls serializable. PiperOrigin-RevId: 542524794