Error raising to stablehlo

[wsmoses]

module {
  func.func private @"##call__Z25gpu__mask_immersed_field_16CompilerMetadataI10StaticSizeI13_180__85__20_E12DynamicCheckvv7NDRangeILi3ES0_I11_12__6__20_ES0_I11_16__16__1_EvvEE11OffsetArrayI7Float64Li3E13CuTracedArrayIS9_Li3ELi1E13_194__99__34_EE5TupleI6Center4FaceSE_E20ImmersedBoundaryGridIS9_8Periodic14RightConnected7Bounded28OrthogonalSphericalShellGridIS9_SI_SJ_SK_28StaticVerticalDiscretizationIS8_IS9_Li1ESA_IS9_Li1ELi1E5_35__EES8_IS9_Li1ESA_IS9_Li1ELi1E5_34__EESO_SQ_ES8_IS9_Li2ESA_IS9_Li2ELi1E9_194__99_EE8TripolarI5Int64SV_SV_EvE16GridFittedBottomI5FieldISE_SE_vvvvS8_IS9_Li3ESA_IS9_Li3ELi1E12_194__99__1_EES9_vvvE23CenterImmersedConditionEvvvES9_#1335$par243"(%arg0: memref<34x99x194xf64, 1>, %arg1: memref<34xf64, 1>, %arg2: memref<1x99x194xf64, 1>) {
    %c-1_i64 = arith.constant -1 : i64
    %c1 = arith.constant 1 : index
    %c1_i64 = arith.constant 1 : i64
    %cst = arith.constant 0.000000e+00 : f64
    affine.parallel (%arg3, %arg4, %arg5) = (0, 0, 0) to (20, 85, 180) {
      %0 = arith.addi %arg4, %c1 : index
      %1 = arith.index_castui %0 : index to i64
      %2 = affine.load %arg1[%arg3 + 7] : memref<34xf64, 1>
      %3 = affine.load %arg2[0, %arg4 + 7, %arg5 + 7] : memref<1x99x194xf64, 1>
      %4 = arith.cmpf ole, %2, %3 {fastmathFlags = #llvm.fastmath<none>} : f64
      %5 = arith.addi %1, %c-1_i64 : i64
      %6 = affine.load %arg2[0, %arg4 + 6, %arg5 + 7] : memref<1x99x194xf64, 1>
      %7 = arith.cmpf ole, %2, %6 {fastmathFlags = #llvm.fastmath<none>} : f64
      %8 = arith.cmpi ult, %5, %c1_i64 : i64
      %9 = arith.ori %8, %7 : i1
      %10 = arith.ori %4, %9 : i1
      %11 = arith.cmpi uge, %5, %c1_i64 : i64
      %12 = arith.andi %11, %10 : i1
      %13 = affine.load %arg0[%arg3 + 7, %arg4 + 7, %arg5 + 7] : memref<34x99x194xf64, 1>
      %14 = arith.select %12, %cst, %13 : f64
      affine.store %14, %arg0[%arg3 + 7, %arg4 + 7, %arg5 + 7] : memref<34x99x194xf64, 1>
    }
    return
  }
}

(base) wmoses@hydra:~/git/Enzyme-JaX$ ./bazel-bin/enzymexlamlir-opt funk.mlir --raise-affine-to-stablehlo
funk.mlir:22:13: error: 'arith.select' op failed to verify that all of {true_value, false_value, result} have same type
      %14 = arith.select %12, %cst, %13 : f64
            ^
funk.mlir:22:13: note: see current operation: %58 = "arith.select"(%56, %57, %55) : (tensor<85x20x180xi1>, tensor<85x20x180xf64>, tensor<20x85x180xf64>) -> tensor<85x20x180xf64>