Incorrect IR generated for Vector Mode AD

```mlir
module {
func.func private @"Const{typeof(rosenbrock)}(Main.rosenbrock)_autodiff"(%arg0: tensor<2xf64>) -> (tensor, tensor<2xf64>) {
%cst = stablehlo.constant dense<1.000000e+02> : tensor
%cst_0 = stablehlo.constant dense<1.000000e+00> : tensor
%0 = stablehlo.slice %arg0 [0:1] : (tensor<2xf64>) -> tensor<1xf64>
%1 = stablehlo.reshape %0 : (tensor<1xf64>) -> tensor
%2 = stablehlo.subtract %cst_0, %1 : tensor
%3 = stablehlo.multiply %2, %2 : tensor
%4 = stablehlo.slice %arg0 [1:2] : (tensor<2xf64>) -> tensor<1xf64>
%5 = stablehlo.reshape %4 : (tensor<1xf64>) -> tensor
%6 = stablehlo.multiply %1, %1 : tensor
%7 = stablehlo.subtract %5, %6 : tensor
%8 = stablehlo.multiply %7, %7 : tensor
%9 = stablehlo.multiply %cst, %8 : tensor
%10 = stablehlo.add %3, %9 : tensor
return %10, %arg0 : tensor, tensor<2xf64>
}
func.func @main(%arg0: tensor<2xf64>, %arg1: tensor<2xf64>, %arg2: tensor<2xf64>) -> (tensor<1xf64>, tensor<1xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) {
%0 = stablehlo.concatenate %arg1, %arg2, dim = 0 : (tensor<2xf64>, tensor<2xf64>) -> tensor<4xf64>
%1 = stablehlo.reshape %0 : (tensor<4xf64>) -> tensor<2x2xf64>
%2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf64>) -> tensor<2x2xf64>
%3:3 = enzyme.fwddiff @"Const{typeof(rosenbrock)}(Main.rosenbrock)_autodiff"(%arg0, %2) {activity = [#enzyme], ret_activity = [#enzyme, #enzyme]} : (tensor<2xf64>, tensor<2x2xf64>) -> (tensor<2xf64>, tensor<2xf64>, tensor<2x2xf64>)
%4 = stablehlo.slice %3#0 [0:1] : (tensor<2xf64>) -> tensor<1xf64>
%5 = stablehlo.slice %3#0 [1:2] : (tensor<2xf64>) -> tensor<1xf64>
return %4, %5, %3#1, %arg1, %arg2 : tensor<1xf64>, tensor<1xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>
}
}
```

```mlir
envs/nested.mlir:40:11: error: the number of elements in start_indices (1) does not match the rank of the operand (0)
%14 = stablehlo.slice %13 [0:1] : (tensor<2xf64>) -> tensor<1xf64>
^
envs/nested.mlir:40:11: error: 'stablehlo.slice' op failed to infer returned types
%14 = stablehlo.slice %13 [0:1] : (tensor<2xf64>) -> tensor<1xf64>
^
envs/nested.mlir:40:11: note: see current operation: %4 = "stablehlo.slice"(%3#0) <{limit_indices = array, start_indices = array, strides = array}> : (tensor) -> tensor<1xf64>
```