These 9 commits are when the Protocol Buffers files have changed:
| 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
The documentation is generated from this commit.
| 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
The documentation is generated from this commit.
| 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