[oryx] harvest after grad

Hi tensorflow probability team!

I was using the oryx library that is maintained in this repo, following an issue from JAX https://github.com/google/jax/issues/3936#issuecomment-668954244.

I am wondering if there is a way to harvest/reap with a grad nested inside, and get the gradient values instead of the primal values. e.g.

import oryx
reap = oryx.core.reap
sow = oryx.core.sow

def f(x): 
  return sow(x ** 3, tag="test", name="x3") + 1

print(reap(grad(f),  tag="test"))
# out: {'x3': DeviceArray(1., dtype=float32)}
# how can we get (something similar to) {'x3': DeviceArray(3., dtype=float32)} ?

Will nest help here?

Thanks!

Start with the Oryx reap/sow and grad example in the issue, then investigate whether nested transformations can expose gradient values rather than primal values. Done means determining whether this behavior is supported and, if not, documenting the required API or implementation scope.