`LLVM ERROR: Failed to infer result type(s)` with sharding

On 633c33a97a3a625804876f1dc40a252e5b88ba4a, the module https://gist.github.com/giordano/debe26b94593d8f228135aab39d24b91 results in
```
loc(callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/packages/LLVM/b3kFs/src/interop/base.jl" in ".">, name = "macro expansion;", linkageName = "macro expansion", file = <"/home/giordano/.julia/packages/LLVM/b3kFs/src/interop/base.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/packages/LLVM/b3kFs/src/interop/base.jl":39:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"none" in ".">, name = "macro expansion;", linkageName = "macro expansion", file = <"none" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["none":0:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"none" in ".">, name = "pointerset;", linkageName = "pointerset", file = <"none" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["none":0:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/packages/LLVM/b3kFs/src/interop/pointer.jl" in ".">, name = "unsafe_store!;", linkageName = "unsafe_store!", file = <"/home/giordano/.julia/packages/LLVM/b3kFs/src/interop/pointer.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/packages/LLVM/b3kFs/src/interop/pointer.jl":88:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, name = "arrayset_bits;", linkageName = "arrayset_bits", file = <"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl":271:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, name = "#arrayset;", linkageName = "#arrayset", file = <"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl":264:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, name = "setindex!;", linkageName = "setindex!", file = <"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl":306:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, name = "setindex!;", linkageName = "setindex!", file = <"/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/dev/Reactant/ext/ReactantCUDAExt.jl":318:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/packages/OffsetArrays/yHW0g/src/OffsetArrays.jl" in ".">, name = "setindex!;", linkageName = "setindex!", file = <"/home/giordano/.julia/packages/OffsetArrays/yHW0g/src/OffsetArrays.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/packages/OffsetArrays/yHW0g/src/OffsetArrays.jl":442:0] at callsite(fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/home/giordano/.julia/packages/KernelAbstractions/sWSE0/src/macros.jl" in ".">, name = "gpu_compute_hydrostatic_free_surface_Gc!;", linkageName = "gpu_compute_hydrostatic_free_surface_Gc!", file = <"/home/giordano/.julia/packages/KernelAbstractions/sWSE0/src/macros.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/home/giordano/.julia/packages/KernelAbstractions/sWSE0/src/macros.jl":322:0] at fused<#llvm.di_subprogram, compileUnit = , sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, name = "gpu_compute_hydrostatic_free_surface_Gc!", linkageName = "julia_gpu_compute_hydrostatic_free_surface_Gc!_66840", file = <"none" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["none":0:0]))))))))))): error: start index 88 is larger than limit index 8 in dimension 2
LLVM ERROR: Failed to infer result type(s):
"stablehlo.slice"(...) {} : (tensor<20x48x96xf64>) -> ( ??? )

[3629415] signal 6 (-6): Aborted
in expression starting at /home/giordano/repo/GB-25/sharding/sharded_baroclinic_instability_simulation_run.jl:72
pthread_kill at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
gsignal at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
abort at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
report_fatal_error at /proc/self/cwd/external/llvm-project/llvm/lib/Support/ErrorHandling.cpp:124
report_fatal_error at /proc/self/cwd/external/llvm-project/llvm/lib/Support/ErrorHandling.cpp:88
reportFatalInferReturnTypesError at /proc/self/cwd/external/llvm-project/mlir/lib/Interfaces/InferTypeOpInterface.cpp:254
build at /proc/self/cwd/bazel-out/k8-opt/bin/external/stablehlo/stablehlo/dialect/StablehloOps.cpp.inc:33791
create, llvm::ArrayRef, llvm::ArrayRef, llvm::ArrayRef > at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/Builders.h:507
matchAndRewrite at /proc/self/cwd/external/enzyme_ad/src/enzyme_ad/jax/Passes/OptimizeCommunication.cpp:791
operator() at /proc/self/cwd/external/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:212 [inlined]
callback_fn<(lambda at external/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:196:9)> at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69 [inlined]
executeAction at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280 [inlined]
matchAndRewrite at /proc/self/cwd/external/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:195
processWorklist at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:615
operator() at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:874 [inlined]
callback_fn<(lambda at external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:873:9)> at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46 [inlined]
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69 [inlined]
executeAction<(anonymous namespace)::GreedyPatternRewriteIteration, long &> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280 [inlined]
simplify at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:872 [inlined]
applyPatternsGreedily at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:919
applyPatternsGreedily at /proc/self/cwd/external/llvm-project/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h:174 [inlined]
applyPatternsAndFoldGreedily at /proc/self/cwd/external/llvm-project/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h:191 [inlined]
runOnOperation at /proc/self/cwd/external/enzyme_ad/src/enzyme_ad/jax/Passes/OptimizeCommunication.cpp:2395
operator() at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:526 [inlined]
callback_fn<(lambda at external/llvm-project/mlir/lib/Pass/Pass.cpp:521:7)> at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46 [inlined]
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69 [inlined]
executeAction at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280 [inlined]
run at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:520
runPipeline at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:592
runPasses at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:905 [inlined]
run at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:885
mlirPassManagerRunOnOp at /proc/self/cwd/external/llvm-project/mlir/lib/CAPI/IR/Pass.cpp:44
mlirPassManagerRunOnOp at /home/giordano/.julia/dev/Reactant/src/mlir/libMLIR_h.jl:8584 [inlined]
run! at /home/giordano/.julia/dev/Reactant/src/mlir/IR/Pass.jl:151
#run_pass_pipeline!#2 at /home/giordano/.julia/dev/Reactant/src/Compiler.jl:874
run_pass_pipeline! at /home/giordano/.julia/dev/Reactant/src/Compiler.jl:869 [inlined]
```