jax

d2653a1e - rewrite axis_index implementation, use custom bind (#2807)

rewrite axis_index implementation, use custom bind (#2807) * rewrite axis_index implementation, use custom bind fixes #2716 Co-authored-by: Trevor Cai <tycai@google.com> * add test for #2716 Co-authored-by: Trevor Cai <tycai@google.com>