[shmap/partial-auto] Fixes lowering for jax.lax.axis_index in shard_map for degenerated shmaps.

[yliu120]

Please take a look at the test I added for the example.

We found that when the manual axis has size 1, the axis_index will be definitely a replicated constant 0 across all the devices. There is no need to add iota and sharding and go full to shard for that, which actually invokes a sharding normalization inside XLA.

Let's say if we do iota and adding the sharding custom ops, then the sharding will be something like,

devices=[1,8,1]<=[8] {replicated, manual} for the iota result.

Then this sharding is normalized to {replicated} in XLA. FullToShard custom call cannot take replicated sharding which will invoke a crash in GSPMD.

[yliu120]

@mattjj @pschuh @yashk2810 Can you please take a look? Thanks.