mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Peter Hawkins 46ac9e2170 Use the default CSR matmul algorithm. 
Previously we requested CUSPARSE_SPMM_CSR_ALG3 in an attempt to get deterministic results from cusparse SpMM CSR matmuls. In the past, Cusparse silently ignored this algorithm choice and used a different algorithm in cases where ALG3 was not supported, but cusparse 12.2.1 removed the silent fallback behavior. Since we're not actually getting deterministic behavior anyway in all cases, use the default algorithm always.

PiperOrigin-RevId: 560867049
2023-08-28 17:49:01 -07:00
..
cpu
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
cuda
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
gpu
Use the default CSR matmul algorithm.
2023-08-28 17:49:01 -07:00
mlir
Move _tpu_ext.cc to jaxlib/mlir/_mlir_libs and set RPATH correctly
2023-07-26 18:25:17 -07:00
mosaic
[Mosaic] Add missing headers in linalg vectorization.
2023-08-28 17:39:47 -07:00
rocm
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
tools
Add generated MLIR _enum_gen.py files to jaxlib wheel.
2023-08-25 19:16:15 +00:00
absl_status_casters.h
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
cpu_feature_guard.c
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ducc_fft.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
gpu_linalg.py
Internal change
2023-03-04 13:24:11 +00:00
gpu_prng.py
[shape_poly] Add support for shape polymorphism for prng GPU custom call
2023-05-02 22:26:58 -07:00
gpu_rnn.py
Resolve unsoundness caught by pytype --strict-none-binding
2023-06-06 06:52:21 -07:00
gpu_solver.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
gpu_sparse.py
Add batching rules to jax.lax.linalg.tridiagonal_solve.
2023-08-10 16:25:59 -07:00
gpu_triton.py
Add get_serialized_metadata function to retrieve metadata from op's opaque data.
2023-06-30 03:23:28 -07:00
handle_pool.h
[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.
2023-05-08 15:37:04 -07:00
hlo_helpers.py
Fix typo in hlo_helpers.py
2023-08-18 16:00:38 +09:00
init.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
jax.bzl
create jax.extend.random as a copy of jax.prng
2023-08-24 14:41:56 -07:00
kernel_helpers.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
kernel_nanobind_helpers.h
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
lapack.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
README.md
Populate long_description for jax & jaxlib
2022-07-13 14:03:32 -07:00
setup.cfg
Export LICENSE.txt file in jaxlib wheels.
2020-12-11 10:20:03 -05:00
setup.py
Enable building jaxlib w/ Mosaic
2023-07-26 03:59:30 -07:00
symlink_files.bzl
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
tpu_mosaic.py
Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic
2023-07-20 18:28:51 -07:00
utils.cc
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00

README.md

jaxlib: support library for JAX

jaxlib is the support library for JAX. While JAX itself is a pure Python package, jaxlib contains the binary (C/C++) parts of the library, including Python bindings, the XLA compiler, the PJRT runtime, and a handful of handwritten kernels. For more information, including installation and build instructions, refer to main JAX README: https://github.com/google/jax/.