jittable pytree flattening/unflattening

[Goodbrake]

I have a problem-specific optimizer that solves a least squares problem at each step using jnp.lstsq. I am currently using ravel_pytree to flatten the gradient into a one dimensional array, and to provide a canonical parameter ordering for the construction of the matrix used in the least squares problem. The problem I face is that the unravel function returned by ravel_pytree is not jit compilable, hence the optimizer's update step cannot be jit compiled. Below is an stripped down version of the optimization step (note this optimizer isn't currently wrapped with @optimizer, so it acts on pytrees directly.)

def update(i, g, state):
  x, unravel = ravel_pytree(state)
  flat_g, _ = ravel_pytree(g)
  step = jnp.linalg.lstsq(matrix(params), flat_g)
  new_x = x - step_size(i) * step
  return unravel(new_x)

Here matrix(params) is a user provided positive semi-definite matrix-valued function of the network parameters of the same form as below.

def matrix(params):
  flat_params, unravel = ravel_pytree(params)

  def flat_vec_fun(flat_ps):
    return vec_fun(unravel(flat_ps))

  jacobian = jacfwd(flat_vec_fun)(flat_params)
  return jnp.matmul(jacobian.T, jacobian)

where here vec_fun is a vector-valued function of the parameters. The matrix actually used in the optimizer is a linear combination of such matrices, and is full rank when the network is minimally parameterized, but I don't think this directly relevant to this feature request.

As you can see, I rely on ravel_pytree to both flatten the parameters into a one dimensional array, and to order the parameters for the construction of the matrix that gets used in the least squares problem. I would like either a version of ravel_pytree that returns a jit compilable unravel function, or a way to make the least squares problem compatible with an existing jit compilable (un)flattening function.

[PhilipVinc]

I have exactly the same use-case as you have in a package I develop.
I ended up using a slightly different version of the ravel function, see it here.

You can quickly check to see if it works for you by simply using nk.jax.tree_ravel instead of the jax-provided function.
If it works for you, you can extract it from there (note: if you are not using complex dtypes, then you can safely replace nkjax.vjp with jax.vjp)

[Goodbrake]

I marked this as the answer, but it actually revealed that my problem was not what I initially thought. The problem was that in my optimizer, the unravel function was passed as part of the optimizer's state, which left the update step unjittable. The problem persisted using your ravel function. Upon refactoring the optimizer to not include the unravel function as part of the state, the jit compilation works with both your ravel, and ravel_pytree.

Thanks anyways!