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
Dan Foreman-Mackey ad1bd38790 Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel. 
This simplifies the lowering logic, and means that we don't get hit with a performance penalty when exporting with shape polymorphism.

PiperOrigin-RevId: 662945116
2024-08-14 09:20:40 -07:00
..
cpu
Determine LAPACK workspaces during SVD kernel runtime
2024-08-13 01:17:44 -07:00
cuda
Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.
2024-08-14 09:20:40 -07:00
gpu
Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.
2024-08-14 09:20:40 -07:00
mlir
[Mosaic TPU][Python] Check validity of VectorLayout on init
2024-08-09 05:28:00 -07:00
mosaic
Integrate LLVM at llvm/llvm-project@c8b5d30f70
2024-08-14 07:09:53 -07:00
rocm
Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.
2024-08-14 09:20:40 -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
Determine LAPACK workspaces during SVD kernel runtime
2024-08-13 01:17:44 -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
Add support for shape polymorphism with lu_pivots_to_permutation.
2024-08-12 03:39:54 -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
Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.
2024-08-14 09:20:40 -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 Singular Value Decomposition to XLA's FFI
2024-08-13 02:41:57 -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/.