Add JAX training tutorial#1302

mtsokol wants to merge 1 commit into

google:maingoogle/grain:mainfrom

mtsokol:jax-training-docsmtsokol/grain:jax-training-docsCopy head branch name to clipboard