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
Tomás Longeri 4452960947 [Mosaic:TPU] In infer ext rule, avoid assigning offsets outside of dst first tile 
Note that offsets outside of first tile are still disabled (for both infer and apply), and once we support it we will want to assign offsets differently, this is mostly to avoid assigning invalid layouts (that may not just be outside the first tile, but outside the vreg slice)

PiperOrigin-RevId: 709168368
2024-12-23 15:49:39 -08:00
..
cpu
Activate Triangular Solve to XLA's FFI
2024-12-11 02:22:37 -08:00
cuda
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
gpu
Port symmetric tridiagonal reduction GPU kernel to FFI.
2024-12-09 12:41:23 -08:00
mlir
Switch mlir bindings from pybind11 to nanobind
2024-12-23 15:10:11 -08:00
mosaic
[Mosaic:TPU] In infer ext rule, avoid assigning offsets outside of dst first tile
2024-12-23 15:49:39 -08:00
rocm
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
tools
[mosaic_gpu] Include Mosaic GPU dialect fiels into jaxlib
2024-12-23 13:46:25 +00: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
[mosaic] Migrated the serialization pass from codegen to pass_boilerplate.h
2024-12-12 14:19:36 -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
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_plugin_extension.cc
Remove cuda include from gpu plugin extension and BUILD
2024-12-11 11:55:51 -06:00
gpu_plugin_extension.h
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08: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
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
Activate Triangular Solve to XLA's FFI
2024-12-11 02:22:37 -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
[mosaic_gpu] Include Mosaic GPU dialect fiels into jaxlib
2024-12-23 13:46:25 +00: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/.