tree_util.tree_vdot bugging on example. about jaxopt HOT 2 CLOSED

import functools

_vdot = functools.partial(jnp.vdot, precision=jax.lax.Precision.HIGHEST)

def _vdot_safe(a, b):
  return _vdot(jnp.asarray(a), jnp.asarray(b))

def tree_vdot(tree_x, tree_y):
  """Compute the inner product <tree_x, tree_y>."""
  vdots = tu.tree_multimap(_vdot_safe, tree_x, tree_y)
  return tu.tree_reduce(operator.add, vdots)