Support forward mode differentiation for SVI

[juanitorduz]

Closes #1726

Trying here a "good first issue".

[fehiepsi]

Looks great to me. Could you add a simple test rwith while loop in the model?

[juanitorduz]

I added a simple test test/infer/test_svi.py::test_forward_mode_differentiation in 532fb5b but it is failing with

    @functools.wraps(update)
    def tree_update(i, grad_tree, opt_state):
      states_flat, tree, subtrees = opt_state
      grad_flat, tree2 = tree_flatten(grad_tree)
      if tree2 != tree:
        msg = ("optimizer update function was passed a gradient tree that did "
               "not match the parameter tree structure with which it was "
               "initialized: parameter tree {} and grad tree {}.")
>       raise TypeError(msg.format(tree, tree2))
E       TypeError: optimizer update function was passed a gradient tree that did not match the parameter tree structure with which it was initialized: parameter tree PyTreeDef({'loc': *, 'scale': *}) and grad tree PyTreeDef(({'loc': *, 'scale': *}, None))

And I am not sure if its because of the implementation of the test is wrong 😑 . Any tips? Thanks

[fehiepsi]

I think you can set has_aux to False.

A faster way, which just does 1 forward pass, is to redefine f:

def wrapper(x):
  out, aux = f(x)
  return out, (out, aux)

grad, (out, aux) = jacfwd(wrapper, has_aux=True)(x)
return (out, aux), grad

We can also apply this trick in hmc forward mode implementation.

[juanitorduz]

ok! So I added your suggestion in ef3a3f7, and now I am getting an error downstream 🤔 . Do I need to unpack somewhere further?

        if progress_bar:
            losses = []
            with tqdm.trange(1, num_steps + 1) as t:
                batch = max(num_steps // 20, 1)
                for i in t:
                    svi_state, loss = jit(body_fn)(svi_state, None)
                    losses.append(jax.device_get(loss))
                    if i % batch == 0:
                        if stable_update:
                            valid_losses = [x for x in losses[i - batch :] if x == x]
                            num_valid = len(valid_losses)
                            if num_valid == 0:
                                avg_loss = float("nan")
                            else:
                                avg_loss = sum(valid_losses) / num_valid
                        else:
>                           avg_loss = sum(losses[i - batch :]) / batch
E                           TypeError: unsupported operand type(s) for +: 'int' and 'dict'

numpyro/infer/svi.py:409: TypeError

[fehiepsi]

sorry, it should be jacfwd of wrapper, not f

[juanitorduz]

It worked in a64c092 !

[juanitorduz]

Same trick in 1259db7 for hmc

[fehiepsi]

Woohoo, thanks for supporting this feature!

[juanitorduz]

Thank you for your guidance @fehiepsi 🙏🙂