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
Qiao Zhang 78963b6020 Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module. 
PiperOrigin-RevId: 490387796
2022-11-22 18:53:29 -08:00
..
blas_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
blas_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
blas.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
BUILD
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00
gpu_kernel_helpers.cc
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00
gpu_kernel_helpers.h
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00
gpu_kernels.cc
Remove unused potrf kernels.
2022-11-17 15:22:13 -08:00
linalg.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
lu_pivot_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07: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
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng_kernels.cu.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
rnn_kernels.cc
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00
rnn_kernels.h
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00
rnn.cc
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00
solver_kernels.cc
Remove unused potrf kernels.
2022-11-17 15:22:13 -08:00
solver_kernels.h
Remove unused potrf kernels.
2022-11-17 15:22:13 -08:00
solver.cc
Remove unused potrf kernels.
2022-11-17 15:22:13 -08:00
sparse_kernels.cc
[sparse] Update the default CUSPARSE SPMV and SPMM algorithms in jaxlib.
2022-11-01 16:01:50 -07:00
sparse_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
sparse.cc
[sparse] Update the default CUSPARSE SPMV and SPMM algorithms in jaxlib.
2022-11-01 16:01:50 -07:00
vendor.h
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-22 18:53:29 -08:00