mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib/gpu
History
Peter Hawkins 95e2d3fc2b [JAX:GPU] Generalize gesvdj kernel to iterate over the unbatched Jacobi kernel in cases that we cannot use the batched kernel. 
If the gesvdj() is preferable to gesvd() absent a batch dimension, even if there is a batch dimension we should prefer a loop of gesvdj() over a loop of gesvd().

PiperOrigin-RevId: 582279549
2023-11-14 04:52:15 -08:00
..
blas_kernels.cc
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
blas_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
blas.cc
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
BUILD
[jax_triton] Add user-specified name field to serialized format.
2023-08-16 02:53:51 -07:00
gpu_kernel_helpers.cc
Fix cupti-related build failure under CUDA 11.
2023-09-27 14:33:30 -07:00
gpu_kernel_helpers.h
[CUDA] Add code to jax initialization that verifies that the CUDA libraries that are found are at least as new as the versions against which JAX was built.
2023-09-27 11:28:40 -07:00
gpu_kernels.cc
[jax_triton] Split C++ only parts of Triton custom callback from Python parts.
2023-07-03 06:52:32 -07:00
linalg.cc
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
lu_pivot_kernels.cc
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
lu_pivot_kernels.cu.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
lu_pivot_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
prng_kernels.cc
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
prng_kernels.cu.cc
Add support for dynamic shapes to GPU threefry2x32 custom call.
2022-12-27 04:48:26 -08:00
prng_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
prng.cc
Switch jaxlib to use nanobind instead of pybind11.
2023-08-24 16:07:56 -07:00
rnn_kernels.cc
Run LSTM test using FP32 math (as opposed to TF32)
2023-09-19 14:45:14 -04:00
rnn_kernels.h
Run LSTM test using FP32 math (as opposed to TF32)
2023-09-19 14:45:14 -04:00
rnn.cc
Run LSTM test using FP32 math (as opposed to TF32)
2023-09-19 14:45:14 -04:00
solver_kernels.cc
[JAX:GPU] Generalize gesvdj kernel to iterate over the unbatched Jacobi kernel in cases that we cannot use the batched kernel.
2023-11-14 04:52:15 -08:00
solver_kernels.h
[shape_poly] Add support for dynamic shapes for eigh
2023-05-02 23:27:59 -07:00
solver.cc
[JAX:GPU] Generalize gesvdj kernel to iterate over the unbatched Jacobi kernel in cases that we cannot use the batched kernel.
2023-11-14 04:52:15 -08:00
sparse_kernels.cc
Add batching rules to jax.lax.linalg.tridiagonal_solve.
2023-08-10 16:25:59 -07:00
sparse_kernels.h
Add batching rules to jax.lax.linalg.tridiagonal_solve.
2023-08-10 16:25:59 -07:00
sparse.cc
Align dummy pointers passed to cusparse to 16 bytes
2023-08-28 12:56:27 -07:00
triton_kernels.cc
Merge pull request #18078 from ROCmSoftwarePlatform:rocm-jax-triton
2023-10-18 11:56:01 -07:00
triton_kernels.h
[ROCm]: Add jax-triton support for ROCm
2023-10-18 07:09:20 +00:00
triton_utils.cc
[jax_triton] Add user-specified name field to serialized format.
2023-08-16 02:53:51 -07:00
triton_utils.h
[jax_triton] Add user-specified name field to serialized format.
2023-08-16 02:53:51 -07:00
triton.cc
[ROCm]: Add jax-triton support for ROCm
2023-10-18 07:09:20 +00:00
triton.proto
[jax_triton] Add user-specified name field to serialized format.
2023-08-16 02:53:51 -07:00
vendor.h
[JAX:GPU] Generalize gesvdj kernel to iterate over the unbatched Jacobi kernel in cases that we cannot use the batched kernel.
2023-11-14 04:52:15 -08:00