mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Jan Naumann e03fe3a06d Implement SVD algorithm based on QR for CPU targets 
In a recent jax release the SvdAlgorithm parameter has been added
to the jax.lax.linalg.svd function. Currently, for CPU targets
still only the divide and conquer algorithm from LAPACK is
supported (gesdd).

This commits adds the functionality to select the QR based
algorithm on CPU as well. Mainly it addes the wrapper code
to call the gesvd function of LAPACK using the FFI interface.

Signed-off-by: Jan Naumann <j.naumann@fu-berlin.de>
2025-02-22 15:24:57 +01:00
..
cpu
Implement SVD algorithm based on QR for CPU targets
2025-02-22 15:24:57 +01:00
cuda
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08:00
gpu
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
mlir
[Mosaic GPU] Implement lowerings for Tile and Transpose transforms from the MLIR dialect.
2025-02-17 01:29:47 -08:00
mosaic
Internal compatibility change
2025-02-21 01:21:56 -08:00
rocm
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08:00
tools
Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.
2025-02-05 10:01:23 -08: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
Migrate JAX MLIR Python dialect extensions to nanobind.
2024-12-13 07:08:28 -08:00
BUILD
[pallas:triton] The lowering now uses PTX instead of Triton IR
2025-02-03 13:21:40 -08: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
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08:00
ffi_helpers.h
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00: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
Fix segfault when old GPU plugins are installed.
2025-02-14 07:26:45 -08:00
gpu_plugin_extension.cc
[pallas:triton] The lowering now uses PTX instead of Triton IR
2025-02-03 13:21:40 -08:00
gpu_plugin_extension.h
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08:00
gpu_prng.py
Fix segfault when old GPU plugins are installed.
2025-02-14 07:26:45 -08:00
gpu_rnn.py
Update RNN kernels to use FFI.
2025-02-06 18:27:58 -08:00
gpu_solver.py
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
gpu_sparse.py
Update the sparse GPU kernels in jaxlib to use the FFI.
2025-02-06 11:45:57 -08: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_python_wheel.bzl
Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.
2025-02-05 10:01:23 -08:00
jax.bzl
Bump the minimum Mac OS X version for x86 builds to 11.0.
2025-02-10 08:51:32 -05:00
kernel_helpers.h
Migrate JAX MLIR Python dialect extensions to nanobind.
2024-12-13 07:08:28 -08: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
Some lax.linalg housekeeping.
2025-02-10 08:27:18 -08:00
pass_boilerplate.h
[mosaic] Migrated the serialization pass from codegen to pass_boilerplate.h
2024-12-12 14:19:36 -08: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
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08:00
setup.py
Increase the minimum SciPy version to 1.11.1.
2025-01-07 16:10:45 -08: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
Add a unittest test extension that runs test cases in parallel using threads.
2025-01-08 09:11:47 -08: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/.