mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Dan Foreman-Mackey 3c014a4c27 Add support for shape polymorphism with lu_pivots_to_permutation. 
This is needed to land support for shape polymorphism with LU decomposition more generally. Most of this change just involves adding the appropriate tests, but I've also updated the "generic" implementation which is used for lowering on CPU to support a dynamic trailing dimension in the input (the `fori_loop` will conditionally lower to a `scan` or `while_loop` as necessary). This change doesn't affect the differentiability (this op doesn't support AD) and the behavior won't change when static shapes are used.

PiperOrigin-RevId: 662024940
2024-08-12 03:39:54 -07:00
..
cpu
Port Eigenvalue Decompositions to XLA's FFI
2024-08-05 03:18:13 -07:00
cuda
Port Getrf GPU kernel from custom call to FFI.
2024-08-01 15:02:25 -07:00
gpu
Merge pull request #22667 from ROCm:rocm-jax-triton-add-get_arch_detail
2024-08-12 02:30:49 -07:00
mlir
[Mosaic TPU][Python] Check validity of VectorLayout on init
2024-08-09 05:28:00 -07:00
mosaic
[Mosaic TPU][Python] Check validity of VectorLayout on init
2024-08-09 05:28:00 -07:00
rocm
[ROCM] Fix BUILD.bazel library source paths
2024-08-07 09:18:20 -05: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
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
Merge pull request #22901 from ROCm:ci_test_harness_vmap
2024-08-06 14:04:57 -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 LU Decomposition to XLA's FFI
2024-08-02 02:22:53 -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/.