`UNIMPLEMENTED: PartitionId instruction is not supported for SPMD partitioning since the meaning is ambiguous`

On d1a2bfaa18ad61645b3bacc5a21925044a236bb4 and with https://github.com/EnzymeAD/Reactant.jl/pull/1171, trying to run
```
XLA_FLAGS="--xla_force_host_platform_device_count=4" julia --project -O0 sharding/sharded_baroclinic_instability_simulation_run.jl
```
in https://github.com/PRONTOLab/GB-25 I get
```
┌ Debug: MLIR module written to /home/giordano/repo/GB-25/sharding/mlir_dumps/2025-04-13T10-55-45.357.2748917/module_298_reactant_first_t..._pre_xla_compile.mlir
└ @ Reactant.MLIR.IR ~/.julia/packages/Reactant/eTubo/src/mlir/IR/Pass.jl:118
┌ Error: Compilation failed, MLIR module written to /home/giordano/repo/GB-25/sharding/mlir_dumps/2025-04-13T10-55-45.357.2748917/module_299_reactant_first_t..._post_xla_compile.mlir
└ @ Reactant.MLIR.IR ~/.julia/packages/Reactant/eTubo/src/mlir/IR/Pass.jl:116
ERROR: LoadError: UNIMPLEMENTED: PartitionId instruction is not supported for SPMD partitioning since the meaning is ambiguous -- whether the instruction is replicated or the data is replicated, and if the latter which data is replicated.

Stacktrace:
[1] reactant_err(msg::Cstring)
@ Reactant.XLA ~/.julia/packages/Reactant/eTubo/src/xla/Utils.jl:12
[2] (::Reactant.XLA.IFRT.var"#33#34"{Bool, Vector{Int64}, Int64, Int64, Bool, Reactant.XLA.IFRT.Client, Reactant.MLIR.IR.Module, Int64})()
@ Reactant.XLA.IFRT ~/.julia/packages/Reactant/eTubo/src/xla/IFRT/LoadedExecutable.jl:89
[3] try_compile_dump_mlir(f::Reactant.XLA.IFRT.var"#33#34"{Bool, Vector{Int64}, Int64, Int64, Bool, Reactant.XLA.IFRT.Client, Reactant.MLIR.IR.Module, Int64}, mod::Reactant.MLIR.IR.Module, pm::Nothing)
@ Reactant.MLIR.IR ~/.julia/packages/Reactant/eTubo/src/mlir/IR/Pass.jl:131
[4] try_compile_dump_mlir
@ ~/.julia/packages/Reactant/eTubo/src/mlir/IR/Pass.jl:126 [inlined]
[5] #compile#32
@ ~/.julia/packages/Reactant/eTubo/src/xla/IFRT/LoadedExecutable.jl:88 [inlined]
```

Here are the two MLIR modules mentioned in the error: https://gist.github.com/giordano/27f6f95d88cb370fc321c43dfa62abde

***Update***: note that this doesn't happen on current [`Reactan#main`](https://github.com/EnzymeAD/Reactant.jl/commit/2daeedecd8004b30166ca02ea7611b4081611a71), so the problem is likely caused by some of the passes enabled by https://github.com/EnzymeAD/Reactant.jl/pull/1171

Reproduce the failure with the Julia command in sharding/sharded_baroclinic_instability_simulation_run.jl using the referenced commit and PR #1171. Compare the two MLIR modules in the linked gist and the behavior of Reactant main at commit 2daeedecd8004b30166ca02ea7611b4081611a71; done means the affected compilation no longer reports the PartitionId SPMD error.

Written by the indexing model from the issue text.