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
Peter Hawkins 5617a02fa4 Remove JAX custom call implementation of batched triangular solve. 
XLA supports batched triangular solve on GPU and has since February 2022, which is older than the minimum jaxlib version. We can therefore delete our implementation and just use XLA's implementation.

PiperOrigin-RevId: 482031830
2022-10-18 15:04:14 -07:00
..
cuda
Remove JAX custom call implementation of batched triangular solve.
2022-10-18 15:04:14 -07:00
mlir
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
rocm
Remove JAX custom call implementation of batched triangular solve.
2022-10-18 15:04:14 -07:00
BUILD
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07: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
cpu_kernels.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ducc_fft_kernels.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ducc_fft_kernels.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ducc_fft.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ducc_fft.fbs
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ducc_fft.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
gpu_linalg.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
gpu_prng.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
gpu_solver.py
Remove JAX custom call implementation of batched triangular solve.
2022-10-18 15:04:14 -07:00
gpu_sparse.py
Reapply: Use input-output aliasing for jaxlib GPU custom calls.
2022-10-10 11:29:18 -07:00
handle_pool.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07: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
Fix bazel build alias
2022-09-26 15:13:12 -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_pybind11_helpers.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
lapack_kernels_using_lapack.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
lapack_kernels.cc
Fix compilation failure in lapack kernel under msan.
2022-10-10 14:24:59 -07:00
lapack_kernels.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
lapack.cc
Avoid loading scipy eagerly.
2022-10-12 19:51:09 +00:00
lapack.py
Avoid loading scipy eagerly.
2022-10-12 19:51:09 +00:00
mhlo_helpers.py
Reapply: Use input-output aliasing for jaxlib GPU custom calls.
2022-10-10 11:29:18 -07:00
README.md
Populate long_description for jax & jaxlib
2022-07-13 14:03:32 -07:00
setup.cfg
Export LICENSE.txt file in jaxlib wheels.
2020-12-11 10:20:03 -05:00
setup.py
Merge pull request #12808 from hawkinsp:py311
2022-10-14 08:56:14 -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

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/.