jax

4a9cc9ff - [Mosaic] Allow passing `ApplyVectorLayoutCtx` to tpu.apply_layout_op.

[Mosaic] Allow passing `ApplyVectorLayoutCtx` to tpu.apply_layout_op. To make it the same with C++ API. While I'm here, fix a bug in test_concatenate. PiperOrigin-RevId: 716016244