Unsupported reduce init for batched reduce

[joelberkeley]

The following batch of reduce

module @root {
  func.func @func1(%arg0: tensor<1xf64>) -> tensor<f64> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f64>) -> tensor<f64>
    %1 = stablehlo.reduce(%arg0 init: %0) applies stablehlo.add across dimensions = [0] : (tensor<1xf64>, tensor<f64>) -> tensor<f64>
    return %1 : tensor<f64>
  }
  func.func @main() -> tensor<2xf64> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<2x1xf64>
    %0 = enzyme.batch @func1(%cst) {batch_shape = array<i64: 2>} : (tensor<2x1xf64>) -> tensor<2xf64>
    return %0 : tensor<2xf64>
  }
}

produces

test.mlir:5:26: error: Unsupported reduce init for batched reduce
    %1 = stablehlo.reduce(%arg0 init: %0) applies stablehlo.add across dimensions = [0] : (tensor<1xf64>, tensor<f64>) -> tensor<f64>
                         ^
test.mlir:5:26: note: see current operation: 
%2 = "stablehlo.reduce"(%arg0, %1) <{dimensions = array<i64: 0>}> ({
^bb0(%arg1: tensor<f64>, %arg2: tensor<f64>):
  %3 = "stablehlo.add"(%arg1, %arg2) : (tensor<f64>, tensor<f64>) -> tensor<f64>
  "stablehlo.return"(%3) : (tensor<f64>) -> ()
}) : (tensor<1xf64>, tensor<f64>) -> tensor<f64>
test.mlir:5:26: error: The shape of reduction-region's argument at index 1 is not compatible with that of reduce-op's input-parameter at index 0
    %1 = stablehlo.reduce(%arg0 init: %0) applies stablehlo.add across dimensions = [0] : (tensor<1xf64>, tensor<f64>) -> tensor<f64>
                         ^