mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Jevin Jiang 9e5edb7015 [Mosaic TPU] Support packed type matmul with arbitrary shapes. 
This cl removes all the shape constrains in matmul for all types.

We only need to mask out subelement on contracting dim. Instead of unpacking data and applying masks, we create a VREG-sized i32 "mask" which contains subelement mask info to logical and with target vreg. Through this way, in order to mask sub-elements, each target vreg only needs to apply 1 op (logical_and) instead of 3 ops (unpacking + select + packing).

PiperOrigin-RevId: 702480077
2024-12-03 14:58:42 -08:00
..
cpu
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
cuda
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
gpu
Integrate Triton up to [9732c047](9732c04701)
2024-12-03 10:55:56 -08:00
mlir
[Mosaic GPU] Add Python bindings for the Mosaic GPU MLIR dialect.
2024-10-31 02:47:30 -07:00
mosaic
[Mosaic TPU] Support packed type matmul with arbitrary shapes.
2024-12-03 14:58:42 -08:00
rocm
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
tools
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -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
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
Remove unneeded dependency from rocm_plugin_extension.
2024-11-21 12:18:11 -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
[ffi] Support handler bundles in GPU plugin extension
2024-09-19 14:51:02 +00: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 some dead code from gpu_prng.py
2024-10-29 09:29:56 -04:00
gpu_rnn.py
Remove support for CUDA 11.
2024-03-25 11:46:39 -07:00
gpu_solver.py
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08: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
Move tsl/platform/{build_config,build_config_root,rules_cc}.bzl to xla/tsl/platform
2024-11-26 15:08:58 -08: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 the FFI implementation of SVD on GPU.
2024-10-17 17:57: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-10 18:00:15 -05:00
setup.py
Mention python 3.13 in docs & package metadata
2024-11-20 08:23:19 -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
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/.