mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib/gpu
History
Chris Jones f238667492 Make JAX-Triton calls serializable. 
PiperOrigin-RevId: 542524794
2023-06-22 04:57:14 -07:00
..
blas_kernels.cc
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
blas_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
blas.cc
Remove an unnecessary include/ from pybind11 include paths.
2022-11-30 14:20:02 -08:00
BUILD
Make JAX-Triton calls serializable.
2023-06-22 04:57:14 -07:00
gpu_kernel_helpers.cc
Add branch prediction to JAX status macros.
2023-05-25 06:23:23 -07:00
gpu_kernel_helpers.h
Add branch prediction to JAX status macros.
2023-05-25 06:23:23 -07:00
gpu_kernels.cc
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
linalg.cc
Remove an unnecessary include/ from pybind11 include paths.
2022-11-30 14:20:02 -08: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
Remove an unnecessary include/ from pybind11 include paths.
2022-11-30 14:20:02 -08:00
rnn_kernels.cc
[GPU] Fix another instance of missing stream synchronization in RNN kernels.
2023-05-09 11:08:24 -07:00
rnn_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
rnn.cc
Remove an unnecessary include/ from pybind11 include paths.
2022-11-30 14:20:02 -08:00
solver_kernels.cc
Use stream-synchronized copy in rnn_kernels.cc.
2023-05-08 13:28:08 -07:00
solver_kernels.h
[shape_poly] Add support for dynamic shapes for eigh
2023-05-02 23:27:59 -07:00
solver.cc
[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.
2023-05-08 15:37:04 -07:00
sparse_kernels.cc
[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.
2023-05-08 15:37:04 -07:00
sparse_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
sparse.cc
[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.
2023-05-08 15:37:04 -07:00
triton.cc
Make JAX-Triton calls serializable.
2023-06-22 04:57:14 -07:00
triton.proto
Make JAX-Triton calls serializable.
2023-06-22 04:57:14 -07:00
vendor.h
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-28 14:31:48 -08:00