/jax

GPU- and TPU-backed NumPy with differentiation and JIT compilation.

Primary LanguagePythonApache License 2.0Apache-2.0

Watchers