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
Dan Foreman-Mackey 9ae1c56c44 Update lu_pivots_to_permutation to use FFI dimensions on GPU. 
The XLA FFI interface provides metadata about buffer dimensions, so quantities
like batch dimensions can be evaluated on the backend, instead of passed as
attributes. This change has the added benefit of allowing this FFI call to
support "vectorized" vmap and dynamic shapes.

PiperOrigin-RevId: 647343656
2024-06-27 09:27:15 -07: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
Move tsl/python to xla/tsl/python
2024-03-29 13:15:21 -07:00
BUILD
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -07:00
cholesky_update_kernel.cc
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -07:00
cholesky_update_kernel.cu.cc
add cholesky changes in bazel
2024-05-18 00:37:09 +00:00
cholesky_update_kernel.h
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -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
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -07:00
linalg.cc
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -07:00
lu_pivot_kernels.cc
Update lu_pivots_to_permutation to use FFI dimensions on GPU.
2024-06-27 09:27:15 -07:00
lu_pivot_kernels.cu.cc
Ported LuPivotsToPermutation to the typed XLA FFI
2024-05-02 08:12:05 -07:00
lu_pivot_kernels.h
Update lu_pivots_to_permutation to use FFI dimensions on GPU.
2024-06-27 09:27:15 -07:00
prng_kernels.cc
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
prng_kernels.cu.cc
Include <cstdint> in files where it is used
2024-03-06 11:58:15 -05:00
prng_kernels.h
Include <cstdint> in files where it is used
2024-03-06 11:58:15 -05: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
Cholesky rank-1 update kernel for JAX.
2024-05-14 15:21:38 -07: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
Move tsl/python to xla/tsl/python
2024-03-29 13:15:21 -07:00
triton_kernels.cc
[Jax/Triton] Skip benchmarking while autotuning for configs that cannot be launched.
2024-04-18 14:35:51 -07:00
triton_kernels.h
[jax_triton] Only use side stream to do autotuning when doing graph capture
2024-02-02 10:48:26 -08: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
[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1
2024-01-19 05:55:41 -08:00
triton.proto
[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1
2024-01-19 05:55:41 -08:00
vendor.h
[ROCm]: Add missing hipStreamWaitEvent API call
2024-03-20 16:58:21 +00:00