https://github.com/google/jax
History
Tip revision: 34b3a7f106baa9099802f246fda6ad3476f32d26 authored by Qiao Zhang on 16 February 2022, 19:26:23 UTC
Allow mesh_util to map a logical axis to 3 physical axes, e.g. for pure data parallel training on TPU.
Tip revision: 34b3a7f
File Mode Size
images
JAX_NeurIPS_2020_demo.ipynb -rw-r--r-- 10.5 KB
JAX_demo.ipynb -rw-r--r-- 19.5 KB
Lorentz_ODE_Solver.ipynb -rw-r--r-- 14.7 KB
Pmap_Cookbook.ipynb -rw-r--r-- 17.9 KB
README.md -rw-r--r-- 4.8 KB
Wave_Equation.ipynb -rw-r--r-- 15.5 KB

README.md

back to top