`Check failed: Product(new_dimensions) == iota_->num_elements() (2 vs. 4)`

MLIR and full XLA dumps at [`2025-04-15T14-52-45.661_K86F.tar.gz`](https://github.com/user-attachments/files/19760570/2025-04-15T14-52-45.661_K86F.tar.gz)

```
[242837] signal 6 (-6): Aborted
in expression starting at /capstor/scratch/cscs/mgiordan/santis/repo/GB-25/sharding/runs/2025-04-15T15-13-51.638_u2br/sharded_baroclinic_instability_simulation_run.jl:67
gsignal at /lib64/libc.so.6 (unknown line)
abort at /lib64/libc.so.6 (unknown line)
warning: address range table at offset 0x17130 has a premature terminator entry at offset 0x17ba0
warning: address range table at offset 0x17130 has a premature terminator entry at offset 0x17bb0
warning: address range table at offset 0x17130 has a premature terminator entry at offset 0x17bc0
[...]
~LogMessageFatal at /proc/self/cwd/external/xla/xla/tsl/platform/default/logging.cc:356
Reshape at /proc/self/cwd/external/xla/xla/hlo/ir/tile_assignment.cc:430
InferShardingFromOperands at /proc/self/cwd/external/xla/xla/service/sharding_propagation.cc:2293
RunToFixPoint at /proc/self/cwd/external/xla/xla/service/sharding_propagation.cc:2976
Run at /proc/self/cwd/external/xla/xla/service/sharding_propagation.cc:3339
RunHelper at /proc/self/cwd/external/xla/xla/hlo/pass/hlo_pass_pipeline.h:150 [inlined]
RunPassesInternal at /proc/self/cwd/external/xla/xla/hlo/pass/hlo_pass_pipeline.cc:198
Run at /proc/self/cwd/external/xla/xla/hlo/pass/hlo_pass_pipeline.cc:338
Run at /proc/self/cwd/external/xla/xla/hlo/pass/hlo_pass_interface.h:85
RunSPMDPasses at /proc/self/cwd/external/xla/xla/service/gpu/gpu_compiler.cc:654 [inlined]
OptimizeHloModule at /proc/self/cwd/external/xla/xla/service/gpu/gpu_compiler.cc:1384
RunHloPasses at /proc/self/cwd/external/xla/xla/service/gpu/gpu_compiler.cc:1845
Compile at /proc/self/cwd/external/xla/xla/service/llvm_compiler.cc:61
BuildExecutables at /proc/self/cwd/external/xla/xla/service/service.cc:285
CompileExecutables at /proc/self/cwd/external/xla/xla/service/local_service.cc:118
Compile at /proc/self/cwd/external/xla/xla/client/local_client.cc:442
CompileInternal at /proc/self/cwd/external/xla/xla/pjrt/pjrt_stream_executor_client.cc:3345
Compile at /proc/self/cwd/external/xla/xla/pjrt/pjrt_stream_executor_client.cc:3438
CompileAndLoad at /proc/self/cwd/external/xla/xla/pjrt/pjrt_stream_executor_client.cc:3452
Create at /proc/self/cwd/external/xla/xla/python/pjrt_ifrt/pjrt_executable.cc:263
Compile at /proc/self/cwd/external/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc:100
ifrt_compile at /proc/self/cwd/API.cpp:1426
#33 at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/xla/IFRT/LoadedExecutable.jl:89
try_compile_dump_mlir at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/mlir/IR/Pass.jl:131
try_compile_dump_mlir at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/mlir/IR/Pass.jl:126 [inlined]
#compile#32 at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/xla/IFRT/LoadedExecutable.jl:88 [inlined]
compile at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/xla/IFRT/LoadedExecutable.jl:74
unknown function (ip: 0x401557870267)
#compile_xla#54 at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/Compiler.jl:2813
compile_xla at /capstor/scratch/cscs/mgiordan/santis/juliaup/depot/packages/Reactant/OHENP/src/Compiler.jl:2765 [inlined]
[...]
```