jax

a67fa87a - [pallas] Generalize support for merging dimensions in fuser reshape.