Autosharding fails with "absl::btree_map::at"

[rahul003]

This is all I see from the logs, don't see full stack trace. I can reproduce this from the autosharding tests as well.
Here's example log from the test

use_sharding_propagation_for_default_shardings: 0
model_resharding_memory_costs: 0
generate_windowed_einsum_strategies: 0
allow_shardings_small_dims_across_many_devices: 0
I0000 00:00:1721332301.859758   16664 auto_sharding.cc:3989] Processing partial mesh shape: [128, 1]
I0000 00:00:1721332301.859811   16664 auto_sharding.cc:4033] Memory consumption lower bound is 0.00232697 GB.
unknown file: Failure
C++ exception with description "absl::btree_map::at" thrown in the test body.
[  FAILED  ] AutoShardingTest.UserShardingTest (3 ms)
[----------] 1 test from AutoShardingTest (3 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test suite ran. (3 ms total)
[  PASSED  ] 0 tests.
[  FAILED  ] 1 test, listed below:
[  FAILED  ] AutoShardingTest.UserShardingTest

[cheshire]

@pratikfegade @mmoffitt could you take a look?