Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.
In most cases, use public jax.* names instead.
"source": [
"@tff.experimental.jax_computation(MODEL_TYPE, BATCH_TYPE)\n",
"def train_on_one_batch(model, batch):\n",
- " grads = jax.api.grad(loss)(model, batch)\n",
+ " grads = jax.grad(loss)(model, batch)\n",
" return collections.OrderedDict([\n",
" (k, model[k] - STEP_SIZE * grads[k]) for k in ['weights', 'bias']\n",
" ])"
