mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-23 02:56:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
407c5c85d3
rocm_jax/jaxlib
History
Ruturaj4 7f9d710015 [ROCm] jaxlib linalg fix
2024-10-16 17:10:35 -05:00
..
cpu
Port Hessenberg Decomposition to XLA's FFI
2024-09-05 01:59:32 -07:00
cuda
Refactor gpusolver kernel definitions into separate build target.
2024-09-12 07:11:36 -07:00
gpu
[ROCM] add missing typename keyword to work with gcc
2024-09-23 14:42:01 -05:00
mlir
Trim StableHLO python binding dependencies
2024-08-28 09:01:15 -07:00
mosaic
[Mosaic TPU] The previous change does not actually force the input offsets read by the rules, but simply disables all the checks. Reverting so that we at least regain the checks until we have a proper fix.
2024-10-01 02:23:52 -07:00
rocm
[ROCm] jaxlib linalg fix
2024-10-16 17:10:35 -05:00
tools
Include the sdy MLIR dialect in jaxlib.
2024-09-27 08:53:31 -07:00
triton
Fix triton capi_objects target to depend on MLIR CAPIIRObjects bazel
2024-09-06 01:06:27 +00:00
absl_status_casters.h
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
[ROCm] Bring up clang support for JAX+XLA
2024-10-16 17:07:55 -05: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
[ffi] Support handler bundles in GPU plugin extension
2024-09-19 14:51:02 +00: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
Simplify logic in jaxlib FFI_ASSIGN_OR_RETURN macro, and fix gcc build.
2024-09-17 11:23:25 -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
Remove forward compatibility checks from cholesky_update lowering.
2024-09-20 08:32:25 -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
Remove forward compatibility checks from lowering of LU decomposition.
2024-09-30 07:23:56 -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
[JAX] bazel build rocm changes
2024-10-16 17:10:02 -05: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 Hessenberg Decomposition to XLA's FFI
2024-10-01 09:20:06 -07:00
README.md
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
rocm_plugin_extension.cc
[JAX] bazel build rocm changes
2024-10-16 17:10:02 -05:00
setup.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -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/jax-ml/jax/.