mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
George Necula 65450d165e Remove forward compatibility mode for old PRGN custom call on GPU 
The backend support for the new custom call was added on June 28th.
Also add backwards compatibility test for the new custom call.

PiperOrigin-RevId: 658011228
2024-07-31 08:10:17 -07:00
..
cpu
Move some common helper functions from lapack_kernels to ffi_helpers.
2024-07-31 07:38:33 -07:00
cuda
Move FFI helper macros from jaxlib/cpu/lapack_kernels.cc to a jaxlib/ffi_helpers.h.
2024-07-10 15:09:45 -07:00
gpu
Remove forward compatibility mode for old PRGN custom call on GPU
2024-07-31 08:10:17 -07:00
mlir
[XLA:Mosaic] Pass rewrite ctx of apply-vector-layout pass to relayout function.
2024-07-23 16:50:45 -07:00
mosaic
[Mosaic TPU] Add support for 1D windows
2024-07-31 05:58:19 -07:00
rocm
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
tools
#sdy Initial set of changes to allow for lowering to the Shardy dialect.
2024-07-23 05:32:06 -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
Move some common helper functions from lapack_kernels to ffi_helpers.
2024-07-31 07:38:33 -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
Replace xla::Status with absl::Status in cuda_plugin_extension.cc.
2024-06-20 13:51:03 -07:00
ducc_fft.py
Add python stub files for jaxlib/cpu C++ Python extensions.
2023-11-28 08:45:24 -08:00
ffi_helpers.h
Move some common helper functions from lapack_kernels to ffi_helpers.
2024-07-31 07:38:33 -07: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
Update lu_pivots_to_permutation to use FFI dimensions on GPU.
2024-06-27 09:27:15 -07:00
gpu_prng.py
Remove forward compatibility mode for old PRGN custom call on GPU
2024-07-31 08:10:17 -07:00
gpu_rnn.py
Remove support for CUDA 11.
2024-03-25 11:46:39 -07:00
gpu_solver.py
[ROCm] Bring-up pjrt support
2024-06-17 16:49:22 +00:00
gpu_sparse.py
[ROCm] Bring-up pjrt support
2024-06-17 16:49:22 +00:00
gpu_triton.py
[ROCm] Bring-up pjrt support
2024-06-17 16:49:22 +00: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
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04: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
Fix gpu_jax_head_jaxlib_pypi_latest job after migrating to plugin structure for jaxlib dependency
2024-07-02 15:35:32 -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
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
lapack.py
Activate Cholesky Factorization Kernel to XLA's FFI
2024-07-25 09:59:28 -07:00
README.md
Populate long_description for jax & jaxlib
2022-07-13 14:03:32 -07:00
rocm_plugin_extension.cc
[ROCM] Replace xla::Status with absl::Status in rocm_plugin_extension.cc
2024-07-02 10:59:53 -05:00
setup.py
Bump minimum SciPy version to 1.10.
2024-07-29 08:50:18 -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/.