mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Yash Katariya 4d8bce1b85 Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. 
```
# Installs Cuda 11 with Cudnn 8.2
$ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

PiperOrigin-RevId: 404134291
2021-10-18 19:56:22 -07:00
..
BUILD
Move cuda .py files to :gpu_support so that if :gpu_support is not present, then internal jaxlib will act like a CPU jaxlib even if --config=cuda is specified.
2021-10-14 13:20:01 -07:00
cpu_feature_guard.c
move #endif so that Windows doesn't have GetXCR0EAX defined twice
2021-06-04 17:26:27 -07:00
cpu_kernels.cc
schur lapack wrapper
2021-09-29 14:29:52 +02:00
cublas_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cublas_kernels.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cublas.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_gpu_kernel_helpers.cc
Fix return value for MakeBatchPointers.
2021-08-09 21:44:46 -07:00
cuda_gpu_kernel_helpers.h
Use absl::Status::message() instead of error_message().
2021-08-09 23:44:36 -07:00
cuda_linalg.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_linalg.py
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_lu_pivot_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_lu_pivot_kernels.cu.cc
Move all Abseil dependencies out of jaxlib CUDA libraries
2021-08-19 16:55:37 -07:00
cuda_lu_pivot_kernels.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_prng_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_prng_kernels.cu.cc
Move all Abseil dependencies out of jaxlib CUDA libraries
2021-08-19 16:55:37 -07:00
cuda_prng_kernels.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_prng.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cuda_prng.py
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusolver_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusolver_kernels.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusolver.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusolver.py
Remove unused backward compatibility code in cusolver.py.
2021-10-18 13:27:10 -04:00
cusparse_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusparse_kernels.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusparse.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
cusparse.py
Move contents of jax.lib to jax._src.lib.
2021-09-23 06:33:55 -07:00
gpu_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
handle_pool.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
init.py
Fix build.py to work on Linux once again.
2020-11-19 13:14:34 -05:00
jax.bzl
Remove jax/BUILD file.
2021-05-19 12:55:34 -07:00
kernel_helpers.h
Use the new "custom call status" facility to report errors in jaxlib
2021-08-09 15:06:39 -07:00
kernel_pybind11_helpers.h
[JAX] Use PocketFFT for FFTs on CPU instead of Eigen.
2020-10-23 14:20:32 -07:00
lapack_kernels.cc
Merge pull request #8033 from SaturdayGenfo:schur-lapack-wrapper
2021-09-29 09:37:49 -07:00
lapack_kernels.h
Merge pull request #8033 from SaturdayGenfo:schur-lapack-wrapper
2021-09-29 09:37:49 -07:00
lapack.cc
Merge pull request #8033 from SaturdayGenfo:schur-lapack-wrapper
2021-09-29 09:37:49 -07:00
lapack.py
schur lapack wrapper
2021-09-29 14:29:52 +02:00
pocketfft_kernels.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
pocketfft_kernels.h
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
pocketfft.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
pocketfft.fbs
[JAX] Use PocketFFT for FFTs on CPU instead of Eigen.
2020-10-23 14:20:32 -07:00
pocketfft.py
Update flatbuffer ver2 EndVector usage.
2021-05-10 16:45:19 -07:00
rocblas.cc
[JAX] Refactor JAX custom kernels to split kernel implementations from Python bindings.
2021-09-02 07:53:09 -07:00
rocm_gpu_kernel_helpers.cc
fix custom_call_status for rocm
2021-08-16 04:39:47 +00:00
rocm_gpu_kernel_helpers.h
Use the new "custom call status" facility to report errors in jaxlib
2021-08-09 15:06:39 -07:00
rocsolver.py
Improve real type conversion in a couple more places.
2021-10-18 13:50:11 -04:00
setup.cfg
Export LICENSE.txt file in jaxlib wheels.
2020-12-11 10:20:03 -05:00
setup.py
Replace _ with - because wheel.py normalizes it to .
2021-10-18 13:47:43 -07:00
version.py
Add a default cuda installation path and more explicit installation paths for CUDA jaxlib.
2021-10-18 19:56:22 -07:00