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
Peter Hawkins 1a193ea189 Fix segfault in cuda_plugin_extension. 
The nanobind switch for the GPU callback code means that we are now using the NumPy APIs rather than pybind11's clone of them. It is important to initialize the NumPy APIs before using them in each module.

PiperOrigin-RevId: 613036056
2024-03-05 18:31:50 -08:00
..
cpu
Merge pull request #19288 from pearu:pearu/int32-overflow
2024-02-20 12:43:16 -08:00
cuda
Update
2024-03-03 19:57:26 +00:00
gpu
[triton] Add clustering support and test
2024-03-04 05:51:10 -08:00
mlir
Don't build the Triton MLIR dialect on Windows
2024-02-21 16:02:54 -08:00
mosaic
[XLA:Mosaic] Use different MXU shape based on the target
2024-03-05 11:14:24 -08:00
rocm
[ROCm]: Add jax-triton support for ROCm
2023-10-18 07:09:20 +00:00
tools
Don't build the Triton MLIR dialect on Windows
2024-02-21 16:02:54 -08:00
triton
Use $(RULEDIR) to avoid an implicit dependency on output_to_genfiles.
2024-02-29 17:40:18 -08:00
absl_status_casters.h
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
Fix segfault in cuda_plugin_extension.
2024-03-05 18:31:50 -08:00
cpu_feature_guard.c
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
cuda_plugin_extension.cc
Fix segfault in cuda_plugin_extension.
2024-03-05 18:31:50 -08:00
ducc_fft.py
Add python stub files for jaxlib/cpu C++ Python extensions.
2023-11-28 08:45:24 -08: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
[PJRT C API] Change build wheel script to build a separate package for cuda kernels.
2023-11-06 09:13:44 -08:00
gpu_prng.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
gpu_rnn.py
Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.
2023-11-17 11:47:14 -08:00
gpu_solver.py
Emit dense arrays for StableHLO ops migrating to dense arrays
2024-01-24 16:41:37 -08:00
gpu_sparse.py
[typing] fix incorrect annotation
2024-01-19 09:10:13 -08:00
gpu_triton.py
[PJRT C API] Change build wheel script to build a separate package for cuda kernels.
2023-11-06 09:13:44 -08: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
Emit dense arrays for StableHLO ops migrating to dense arrays
2024-01-24 16:41:37 -08: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
Fixed internal CI builds
2024-02-20 02:42:14 -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 Windows build failure.
2023-09-05 10:14:26 -04:00
lapack.py
Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.
2023-11-17 11:47:14 -08:00
README.md
Populate long_description for jax & jaxlib
2022-07-13 14:03:32 -07:00
setup.py
Include Triton files into the jaxlib wheel
2024-01-16 15:28:12 +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/google/jax/.