https://github.com/google/jax
History
Tip revision: 5c4e88d8799d8e24ca5bfba5a84c177a66b21f92 authored by Yash Katariya on 09 August 2023, 02:00:31 UTC
[Memories] Allow device_put outside jax.jit to work with numpy arrays, numpy scalars, python scalars and python ints with memory kinds
Tip revision: 5c4e88d
File Mode Size
lorentz.png -rw-r--r-- 560.6 KB
nested_pmap.png -rw-r--r-- 23.3 KB
wave_movie.gif -rw-r--r-- 5.6 MB

back to top