jax

ec89afbc - [Mosaic SC] Generalize broadcastSubelements to support SC shapes

[Mosaic SC] Generalize broadcastSubelements to support SC shapes We only needed to stop hard-coding the target shape to be 2D. We use the same convention for packed dim as for TC, so we can reuse the same helper. PiperOrigin-RevId: 925309748