def _adasum_merge(g1, g2):
  """Adasum gradient composition, see https://arxiv.org/pdf/2006.02924.pdf."""
  frac1 = jnp.vdot(g1, g2) / (2 * jnp.vdot(g1, g1) + 1e-30)
  frac2 = jnp.vdot(g1, g2) / (2 * jnp.vdot(g2, g2) + 1e-30)
  return  (1 - frac1) * g1 + (1 - frac2) * g2
Adasum gradient composition, see https://arxiv.org/pdf/2006.02924.pdf.
trax/optimizers/trainer.py