Yash Katariya 06ad3528e9 Use _make_lengths_same for explicit mode too. ...

We add `None`'s when ndim > len(sharding.spec) and only remove `None`s when `ndim < len(sharding.spec)`. If sharded axes exist, then we error out when removing specs.

PiperOrigin-RevId: 748735303

2025-04-17 10:48:46 -07:00

..

__init__.py

[Jax][Pallas][Mosaic] Implement platform dependent diag, with branch selection driven by constant prop in mosaic lowering.

2024-12-22 00:50:51 -08:00

common.py

Fix dtype failures in JaxGroupedQueryAttentionReferenceTest.

2025-04-10 11:04:43 -07:00

conditionals.py

Fix dtype failures in JaxGroupedQueryAttentionReferenceTest.

2025-04-10 11:04:43 -07:00

for_loop.py

[better_errors] Make it explicit that debug_info is not None.

2025-02-13 22:07:04 -08:00

loops.py

Use _make_lengths_same for explicit mode too.

2025-04-17 10:48:46 -07:00

solves.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00