xla

f7e27bea - add customized ring in jax