Inner Loop Gradients

[pharringtonp19]

In MAML, I would like to be able to specify which layers to update during the inner loop (i.e. compute and apply the gradient via SGD to only a subset of the parameters denoted by p) -- I am using Stax. What would be the cleanest way to do this? -- Thanks

def inner_update(p, x, y):
  grads = grad(loss)(p, x, y)                  # Compute the gradient
  return tree_multimap(inner_sgd_fn, grads, p) # Apply gradient update

Haiku offers a clean solution