mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Kevin Gleason 78d5b75b0d Trim StableHLO python binding dependencies 
With proper CAPI in place these dependencies are no longer needed, llvm support needed for string ostream for string APIs.

PiperOrigin-RevId: 668476145
2024-08-28 09:01:15 -07:00
..
cpu
Clean up BUILD files.
2024-08-26 09:11:17 -07:00
cuda
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
gpu
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
mlir
Trim StableHLO python binding dependencies
2024-08-28 09:01:15 -07:00
mosaic
Clean up BUILD files.
2024-08-26 09:11:17 -07:00
rocm
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -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
Clean up BUILD files.
2024-08-26 09:11:17 -07:00
absl_status_casters.h
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
Clean up BUILD files.
2024-08-26 09:11:17 -07:00
cpu_feature_guard.c
Added PyUnstable_Module_SetGIL to PyInit_cpu_feature_guard
2024-08-20 02:32:37 +02: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
Port QR factorization GPU kernel to FFI.
2024-08-20 05:07:04 -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
Port the GPU Cholesky update custom call to the FFI.
2024-08-20 05:46:03 -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
Activate GPU kernel for LU decomposition.
2024-08-21 05:08:41 -07:00
gpu_sparse.py
[ROCm] Bring-up pjrt support
2024-06-17 16:49:22 +00:00
gpu_triton.py
[ROCm]: Add get_arch_details for triton kernel call
2024-08-12 06:16:27 +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
Deprecate using build_cuda_plugin_from_source flag and rely on jaxlib_build config.
2024-08-08 11:58:13 -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 Eigenvalue Decompositions to XLA's FFI
2024-08-28 03:53:49 -07:00
README.md
Populate long_description for jax & jaxlib
2022-07-13 14:03:32 -07:00
rocm_plugin_extension.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07: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/.