mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
jax authors 71c19b779d Rewrite vector.contraction with bf16 accumulator and output into a 
contraction with f32 accumulator and output, where the accumulator is
extended and the output truncated. For targets that do not support bf16
matmul, the lhs and rhs are extended to f32.

PiperOrigin-RevId: 642051952
2024-06-10 16:02:46 -07:00
..
cpu
Refactor kernel function assigment
2024-06-07 08:20:31 -07:00
cuda
Add initialization annotations (for the benefit of MSAN) to variables that are initialized by external functions.
2024-06-10 06:21:16 -07:00
gpu
add cholesky changes in bazel
2024-05-18 00:37:09 +00:00
mlir
[Mosaic GPU] Integrate the ExecutionEngine with the jaxlib GPU plugin
2024-05-30 01:46:23 -07:00
mosaic
Rewrite vector.contraction with bf16 accumulator and output into a
2024-06-10 16:02:46 -07:00
rocm
add cholesky changes in bazel
2024-05-18 00:37:09 +00:00
tools
[Mosaic GPU] Integrate the ExecutionEngine with the jaxlib GPU plugin
2024-05-30 01:46:23 -07:00
triton
Removed inspect.signature() call from jaxlib.triton.dialect.ScanOp
2024-03-11 13:30:41 -07:00
absl_status_casters.h
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
Expose XLA FFI headers to bazel build and re-enable tests
2024-06-04 10:14:43 -07:00
cpu_feature_guard.c
Add a comment about x86 builds of Python to cpu_feature_guard
2024-05-29 15:25:44 -04:00
cuda_plugin_extension.cc
[xla:python] Add support for registering custom call targets for all XLA execution stages and for XLA FFI traits
2024-05-24 10:34:46 -07:00
ducc_fft.py
Add python stub files for jaxlib/cpu C++ Python extensions.
2023-11-28 08:45:24 -08:00
gpu_common_utils.py
Improve the gpu lowering error message if users forget link the gpu library.
2023-09-11 16:14:18 -07:00
gpu_linalg.py
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -07:00
gpu_prng.py
Remove support for CUDA 11.
2024-03-25 11:46:39 -07:00
gpu_rnn.py
Remove support for CUDA 11.
2024-03-25 11:46:39 -07:00
gpu_solver.py
Clean up version switches from dense array migration
2024-05-28 10:58:51 -07:00
gpu_sparse.py
Remove support for CUDA 11.
2024-03-25 11:46:39 -07:00
gpu_triton.py
Remove support for CUDA 11.
2024-03-25 11:46:39 -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
Clean up version switches from dense array migration
2024-05-28 10:58:51 -07: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
Added filelock to py_deps
2024-06-07 13:16:33 -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
Move tsl/python to xla/tsl/python
2024-03-29 13:15:21 -07:00
lapack.py
Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.
2023-11-17 11:47:14 -08:00
README.md
Populate long_description for jax & jaxlib
2022-07-13 14:03:32 -07:00
setup.py
Update cuDNN to version 9.1.1 in JAX
2024-05-24 10:10:21 -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
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/.