Could LBFGS work with pytrees?

Similar to this previous issue (https://github.com/blackjax-devs/blackjax/issues/214), would it be possible to modify minimize_lbfgs to also work with the initial guess x0 being something more flexible like a tuple of arrays? That would be really handy so that I could use the same logprob function to draw samples and to find the (local) max of the posterior.

[ahwillia]

Happy to give it a shot. I am somewhat new to jax just learned that jax.flatten_util.ravel_pytree exists. I’m not sure I understand the cases where vmap(unravel_fn)(...) is needed though.

Is the following sketch what you two are thinking?

def minimize_lbfgs(func, x0, **kwargs):
    if isinstance(x0, jax.numpy.DeviceArray):
        return _minimize_lbfgs(func, x0, **kwargs)
    else:
        x0_raveled, unravel_fn = ravel_pytree(init_params)
        return _minimize_lbfgs(lambda x: func(unravel_fn(x)), x0_raveled, **kwargs)

Where minimize_lbfgs -> _minimize_lfbgs.

[rlouf]

Closing as the corresponding PR was merged. Thank you!