Files

 main

/

gpu

/

Copy path

Directory actions

More options

Directory actions

More options

Latest commit

gnecula

and

jax authors

Add support for dynamic shapes to GPU threefry2x32 custom call.

Dec 27, 2022

7d452ad · Dec 27, 2022

History

History

This branch is 11821 commits behind jax-ml/jax:main.

/

gpu

/

Top

Folders and files

Name		Name	Last commit message	Last commit date
parent directory ..
BUILD		BUILD	Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.	Nov 28, 2022
blas.cc		blas.cc	Remove an unnecessary include/ from pybind11 include paths.	Nov 30, 2022
blas_kernels.cc		blas_kernels.cc	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
blas_kernels.h		blas_kernels.h	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
gpu_kernel_helpers.cc		gpu_kernel_helpers.cc	Rollback of "Add CUDNN custom call for LSTM. Exposed as jax.experimen…	Nov 23, 2022
gpu_kernel_helpers.h		gpu_kernel_helpers.h	Rollback of "Add CUDNN custom call for LSTM. Exposed as jax.experimen…	Nov 23, 2022
gpu_kernels.cc		gpu_kernels.cc	Remove unused potrf kernels.	Nov 17, 2022
linalg.cc		linalg.cc	Remove an unnecessary include/ from pybind11 include paths.	Nov 30, 2022
lu_pivot_kernels.cc		lu_pivot_kernels.cc	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
lu_pivot_kernels.cu.cc		lu_pivot_kernels.cu.cc	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
lu_pivot_kernels.h		lu_pivot_kernels.h	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
prng.cc		prng.cc	Remove an unnecessary include/ from pybind11 include paths.	Nov 30, 2022
prng_kernels.cc		prng_kernels.cc	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
prng_kernels.cu.cc		prng_kernels.cu.cc	Add support for dynamic shapes to GPU threefry2x32 custom call.	Dec 27, 2022
prng_kernels.h		prng_kernels.h	Add support for dynamic shapes to GPU threefry2x32 custom call.	Dec 27, 2022
rnn.cc		rnn.cc	Remove an unnecessary include/ from pybind11 include paths.	Nov 30, 2022
rnn_kernels.cc		rnn_kernels.cc	Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.	Nov 28, 2022
rnn_kernels.h		rnn_kernels.h	Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.	Nov 28, 2022
solver.cc		solver.cc	Remove an unnecessary include/ from pybind11 include paths.	Nov 30, 2022
solver_kernels.cc		solver_kernels.cc	Remove unused potrf kernels.	Nov 17, 2022
solver_kernels.h		solver_kernels.h	Remove unused potrf kernels.	Nov 17, 2022
sparse.cc		sparse.cc	Remove an unnecessary include/ from pybind11 include paths.	Nov 30, 2022
sparse_kernels.cc		sparse_kernels.cc	[sparse] Update the default CUSPARSE SPMV and SPMM algorithms in jaxlib.	Nov 1, 2022
sparse_kernels.h		sparse_kernels.h	Merge CUDA and ROCM kernel code in jaxlib.	Oct 25, 2022
vendor.h		vendor.h	Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.	Nov 28, 2022
View all files