From d4b8754392f220108cf4351554aa3feb73b1c4c5 Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion@google.com>
Date: Fri, 2 Feb 2024 10:25:55 -0800
Subject: [PATCH] Disallow unranked dynamism in ops that take output shape as
 an operand

This includes:
- dynamic_broadcast_in_dim
- dynamic_iota
- dynamic_reshape

https://github.com/openxla/stablehlo/pull/1881
---
 .../linalg/tests/miscellaneous.mlir           |  14 ++--
 stablehlo/dialect/Base.td                     |   5 ++
 stablehlo/dialect/StablehloOps.td             |  12 ++--
 stablehlo/dialect/TypeInference.cpp           |  12 +---
 stablehlo/tests/ops_stablehlo.mlir            |  61 +++++++++++-------
 stablehlo/tests/stablehlo_refine_shapes.mlir  |  18 +++---
 .../stablehlo_legalize_to_vhlo.0_10_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_10_0.mlir.bc | Bin 20214 -> 20214 bytes
 .../stablehlo_legalize_to_vhlo.0_11_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_11_0.mlir.bc | Bin 20299 -> 20299 bytes
 .../stablehlo_legalize_to_vhlo.0_12_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_12_0.mlir.bc | Bin 20943 -> 20943 bytes
 .../stablehlo_legalize_to_vhlo.0_13_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_13_0.mlir.bc | Bin 20943 -> 20943 bytes
 .../stablehlo_legalize_to_vhlo.0_14_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_14_0.mlir.bc | Bin 20593 -> 20593 bytes
 .../stablehlo_legalize_to_vhlo.0_15_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_15_0.mlir.bc | Bin 16618 -> 16618 bytes
 .../stablehlo_legalize_to_vhlo.0_16_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_16_0.mlir.bc | Bin 16746 -> 16746 bytes
 .../stablehlo_legalize_to_vhlo.0_17_0.mlir    |   6 +-
 .../stablehlo_legalize_to_vhlo.0_17_0.mlir.bc | Bin 17638 -> 17636 bytes
 .../stablehlo_legalize_to_vhlo.0_9_0.mlir     |   6 +-
 .../stablehlo_legalize_to_vhlo.0_9_0.mlir.bc  | Bin 20049 -> 20049 bytes
 .../vhlo/stablehlo_legalize_to_vhlo.mlir      |   6 +-
 25 files changed, 99 insertions(+), 83 deletions(-)

diff --git a/stablehlo/conversions/linalg/tests/miscellaneous.mlir b/stablehlo/conversions/linalg/tests/miscellaneous.mlir
index 5ec4206adcb..643c109e87f 100644
--- a/stablehlo/conversions/linalg/tests/miscellaneous.mlir
+++ b/stablehlo/conversions/linalg/tests/miscellaneous.mlir
@@ -601,9 +601,9 @@ func.func @iota_complexf32() -> tensor<7x10xcomplex<f32>> {
 
 // CHECK: #[[RESULT_MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK: func @dynamic_iota_f32
-// CHECK-SAME: %[[SHAPE:.*]]: tensor<?xi32>
-func.func @dynamic_iota_f32(%shape: tensor<?xi32>) -> tensor<?x?x8xf32> {
-  %result = "stablehlo.dynamic_iota"(%shape) {iota_dimension = 1 : i64} : (tensor<?xi32>) -> (tensor<?x?x8xf32>)
+// CHECK-SAME: %[[SHAPE:.*]]: tensor<3xi32>
+func.func @dynamic_iota_f32(%shape: tensor<3xi32>) -> tensor<?x?x8xf32> {
+  %result = "stablehlo.dynamic_iota"(%shape) {iota_dimension = 1 : i64} : (tensor<3xi32>) -> (tensor<?x?x8xf32>)
   func.return %result : tensor<?x?x8xf32>
 }
 // CHECK: %[[V1:.*]] = tensor.extract %[[SHAPE]][%c0]
@@ -622,10 +622,10 @@ func.func @dynamic_iota_f32(%shape: tensor<?xi32>) -> tensor<?x?x8xf32> {
 // -----
 
 // CHECK: #[[RESULT_MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-// CHECK: func @dyanmic_iota_ui32
-// CHECK-SAME: %[[SHAPE:.*]]: tensor<?xi32>
-func.func @dyanmic_iota_ui32(%shape: tensor<?xi32>) -> tensor<?x?x8xui32> {
-  %result = "stablehlo.dynamic_iota"(%shape) {iota_dimension = 1 : i64} : (tensor<?xi32>) -> (tensor<?x?x8xui32>)
+// CHECK: func @dynamic_iota_ui32
+// CHECK-SAME: %[[SHAPE:.*]]: tensor<3xi32>
+func.func @dynamic_iota_ui32(%shape: tensor<3xi32>) -> tensor<?x?x8xui32> {
+  %result = "stablehlo.dynamic_iota"(%shape) {iota_dimension = 1 : i64} : (tensor<3xi32>) -> (tensor<?x?x8xui32>)
   func.return %result : tensor<?x?x8xui32>
 }
 // CHECK: %[[V1:.*]] = tensor.extract %[[SHAPE]][%c0]
diff --git a/stablehlo/dialect/Base.td b/stablehlo/dialect/Base.td
index efcede4ab33..51ca0b8c001 100644
--- a/stablehlo/dialect/Base.td
+++ b/stablehlo/dialect/Base.td
@@ -117,6 +117,8 @@ def HLO_PredTensor : TensorOf<[HLO_Pred]>;
 
 def HLO_Tensor : TensorOf<[HLO_Float, HLO_Pred, HLO_Int, HLO_Complex, HLO_QuantizedInt]>;
 
+def HLO_RankedTensor : RankedTensorOf<[HLO_Float, HLO_Pred, HLO_Int, HLO_Complex, HLO_QuantizedInt]>;
+
 def HLO_ComplexTensor : TensorOf<[HLO_Complex]>;
 
 def HLO_Tuple : NestedTupleOf<[HLO_Tensor, HLO_Token]>;
@@ -130,6 +132,9 @@ def HLO_DimensionValue : AnyTypeOf<[Index, HLO_Int]>;
 // Dynamic representation of a shape vector as a tensor.
 def HLO_DimensionTensor : 1DTensorOf<[HLO_DimensionValue]>;
 
+// Static representation of a shape vector as a tensor.
+def HLO_StaticDimensionTensor : RankedTensorOf<[HLO_DimensionValue], [HasStaticShapePred, HasAnyRankOfPred<[1]>], "statically shaped 1-dimensional tensor">;
+
 //===----------------------------------------------------------------------===//
 // HLO combined type definitions.
 //===----------------------------------------------------------------------===//
diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td
index 2ca6afeb9eb..8c5c7dc7369 100644
--- a/stablehlo/dialect/StablehloOps.td
+++ b/stablehlo/dialect/StablehloOps.td
@@ -141,8 +141,8 @@ def StableHLO_DynamicIotaOp: StableHLO_ShapedInterfaceOp<"dynamic_iota", [Pure]>
     ```
   }];
 
-  let arguments = (ins HLO_DimensionTensor:$output_shape, I64Attr:$iota_dimension);
-  let results = (outs HLO_Tensor:$result);
+  let arguments = (ins HLO_StaticDimensionTensor:$output_shape, I64Attr:$iota_dimension);
+  let results = (outs HLO_RankedTensor:$result);
   let hasVerifier = 1;
 
   let assemblyFormat = [{
@@ -1910,13 +1910,13 @@ def StableHLO_DynamicBroadcastInDimOp : StableHLO_ShapedInterfaceOp<
   }];
   let arguments = (ins
     HLO_Tensor:$operand,
-    HLO_DimensionTensor:$output_dimensions,
+    HLO_StaticDimensionTensor:$output_dimensions,
     DenseI64ArrayAttr:$broadcast_dimensions,
     OptionalAttr<DenseI64ArrayAttr>:$known_expanding_dimensions,
     OptionalAttr<DenseI64ArrayAttr>:$known_nonexpanding_dimensions
   );
 
-  let results = (outs HLO_Tensor);
+  let results = (outs HLO_RankedTensor);
 
   let builders = [
       OpBuilder<(ins
@@ -2508,8 +2508,8 @@ def StableHLO_DynamicReshapeOp: StableHLO_ShapedInterfaceOp<"dynamic_reshape", [
     ```
   }];
 
-  let arguments = (ins HLO_Tensor:$operand, HLO_DimensionTensor:$output_shape);
-  let results = (outs HLO_Tensor:$result);
+  let arguments = (ins HLO_Tensor:$operand, HLO_StaticDimensionTensor:$output_shape);
+  let results = (outs HLO_RankedTensor:$result);
 
   let hasVerifier = 1;
 
diff --git a/stablehlo/dialect/TypeInference.cpp b/stablehlo/dialect/TypeInference.cpp
index 159cb5e275c..811980be302 100644
--- a/stablehlo/dialect/TypeInference.cpp
+++ b/stablehlo/dialect/TypeInference.cpp
@@ -3470,11 +3470,7 @@ LogicalResult verifyDynamicBroadcastInDimOp(
     std::optional<ArrayRef<int64_t>> knownNonexpandingDimensions,
     Value result) {
   auto operandType = operand.getType().dyn_cast<RankedTensorType>();
-  auto resultType = result.getType().dyn_cast<RankedTensorType>();
-
-  // If either the operand or result are unranked, there is very little
-  // to verify statically.
-  if (!operandType || !resultType) return success();
+  auto resultType = result.getType().cast<RankedTensorType>();
 
   auto outputDimensionsType =
       outputDimensions.getType().cast<RankedTensorType>();
@@ -3555,13 +3551,12 @@ LogicalResult verifyDynamicIotaOp(std::optional<Location> location,
                                   Value outputShape, int64_t iotaDimension,
                                   Value result) {
   auto shape = result.getType().cast<ShapedType>();
+
   if (!isCompatibleForHloTypeInference(outputShape, shape))
     return emitOptionalError(
         location, "output_shape is incompatible with return type of operation ",
         result.getType());
 
-  if (!shape.hasRank()) return success();
-
   if (iotaDimension >= shape.getRank() || iotaDimension < 0)
     return emitOptionalError(
         location,
@@ -3616,8 +3611,7 @@ LogicalResult verifyDynamicReshapeOp(std::optional<Location> location,
                                      Value outputShape, Value result) {
   auto resultType = result.getType().cast<ShapedType>();
   auto outputShapeType = outputShape.getType().cast<ShapedType>();
-  if (resultType.hasRank() && outputShapeType.hasStaticShape() &&
-      outputShapeType.getDimSize(0) != resultType.getRank())
+  if (outputShapeType.getDimSize(0) != resultType.getRank())
     return emitOptionalError(location,
                              "output should have a rank equal to the number of "
                              "elements in output_shape");
diff --git a/stablehlo/tests/ops_stablehlo.mlir b/stablehlo/tests/ops_stablehlo.mlir
index ad4f6a03265..22060062947 100644
--- a/stablehlo/tests/ops_stablehlo.mlir
+++ b/stablehlo/tests/ops_stablehlo.mlir
@@ -1010,15 +1010,6 @@ func.func @dynamic_broadcast_in_dim(%arg0: tensor<?x?xi32>, %shape: tensor<3xi64
   %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %shape) {broadcast_dimensions = array<i64: 1, 2>} : (tensor<?x?xi32>, tensor<3xi64>) -> tensor<?x?x?xi32>
   func.return %0 : tensor<?x?x?xi32>
 }
-
-// -----
-
-// CHECK-LABEL: func @dynamic_broadcast_in_dim_unranked
-func.func @dynamic_broadcast_in_dim_unranked(%arg0: tensor<?x?xi32>, %shape: tensor<3xi64>) -> tensor<*xi32> {
-  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %shape) {broadcast_dimensions = array<i64: 1, 2>} : (tensor<?x?xi32>, tensor<3xi64>) -> tensor<*xi32>
-  func.return %0 : tensor<*xi32>
-}
-
 // -----
 
 // CHECK-LABEL: func @dynamic_broadcast_in_dim_unknown_dim
@@ -1079,6 +1070,22 @@ func.func @dynamic_broadcast_in_dim_too_large(%arg0: tensor<1xf32>, %shape: tens
 
 // -----
 
+func.func @dynamic_broadcast_in_dim_unranked_result(%arg0: tensor<?x?xi32>, %shape: tensor<3xi64>) -> tensor<*xi32> {
+  // expected-error@+1 {{op result #0 must be ranked tensor}}
+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %shape) {broadcast_dimensions = array<i64: 1, 2>} : (tensor<?x?xi32>, tensor<3xi64>) -> tensor<*xi32>
+  func.return %0 : tensor<*xi32>
+}
+
+// -----
+
+func.func @dynamic_broadcast_in_dim_dynamic_output_shape(%arg0: tensor<?x?xi32>, %shape: tensor<?xi64>) -> tensor<7x8x9xi32> {
+  // expected-error@+1 {{op operand #1 must be statically shaped}}
+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %shape) {broadcast_dimensions = array<i64: 1, 2>} : (tensor<?x?xi32>, tensor<?xi64>) -> tensor<7x8x9xi32>
+  func.return %0 : tensor<7x8x9xi32>
+}
+
+// -----
+
 // CHECK-LABEL: func @broadcast_in_dim
 func.func @broadcast_in_dim(%arg0: tensor<1x2xi32>) -> tensor<1x2x2xi32> {
   %0 = "stablehlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = array<i64: 1, 2>} : (tensor<1x2xi32>) -> tensor<1x2x2xi32>
@@ -3391,13 +3398,6 @@ func.func @dynamic_reshape(%arg0: tensor<?xf32>, %shape: tensor<2xindex>) -> ten
 
 // -----
 
-func.func @dynamic_reshape_unranked(%arg0: tensor<?xf32>, %shape: tensor<2xindex>) -> tensor<*xf32> {
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %shape) : (tensor<?xf32>, tensor<2xindex>) -> tensor<*xf32>
-  func.return %0 : tensor<*xf32>
-}
-
-// -----
-
 func.func @dynamic_reshape_incompatible_shapes(%arg0: tensor<?xf32>, %shape: tensor<2xindex>) -> tensor<?xf32> {
   // expected-error @+1 {{output should have a rank equal to the number of elements in output_shape}}
   %0 = "stablehlo.dynamic_reshape"(%arg0, %shape) : (tensor<?xf32>, tensor<2xindex>) -> tensor<?xf32>
@@ -3407,7 +3407,7 @@ func.func @dynamic_reshape_incompatible_shapes(%arg0: tensor<?xf32>, %shape: ten
 // -----
 
 func.func @dynamic_reshape_output_shape_negative_size(%arg0: tensor<4xf32>) -> tensor<1x4xf32> {
-  // @expected-error@+2 {{output_shape is incompatible with return type of operation 'tensor<1x4xf32>'}}
+  // expected-error@+2 {{output_shape is incompatible with return type of operation 'tensor<1x4xf32>'}}
   %0 = stablehlo.constant dense<[-1, 1]> : tensor<2xi64>
   %1 = stablehlo.dynamic_reshape %arg0, %0 : (tensor<4xf32>, tensor<2xi64>) -> tensor<1x4xf32>
   return %1 : tensor<1x4xf32>
@@ -3416,7 +3416,7 @@ func.func @dynamic_reshape_output_shape_negative_size(%arg0: tensor<4xf32>) -> t
 // -----
 
 func.func @dynamic_reshape_output_shape_mismatching_size(%arg0: tensor<4xf32>) -> tensor<1x4xf32> {
-  // @expected-error@+2 {{output_shape is incompatible with return type of operation 'tensor<1x4xf32>'}}
+  // expected-error@+2 {{output_shape is incompatible with return type of operation 'tensor<1x4xf32>'}}
   %0 = stablehlo.constant dense<[1, 1]> : tensor<2xi64>
   %1 = stablehlo.dynamic_reshape %arg0, %0 : (tensor<4xf32>, tensor<2xi64>) -> tensor<1x4xf32>
   return %1 : tensor<1x4xf32>
@@ -3424,6 +3424,22 @@ func.func @dynamic_reshape_output_shape_mismatching_size(%arg0: tensor<4xf32>) -
 
 // -----
 
+func.func @dynamic_reshape_unranked_result(%arg0: tensor<?xf32>, %shape: tensor<2xindex>) -> tensor<*xf32> {
+  // expected-error@+1 {{op result #0 must be ranked tensor}}
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %shape) : (tensor<?xf32>, tensor<2xindex>) -> tensor<*xf32>
+  func.return %0 : tensor<*xf32>
+}
+
+// -----
+
+func.func @dynamic_reshape_dynamic_output_shape(%arg0: tensor<?xf32>, %shape: tensor<?xindex>) -> tensor<1x4xf32> {
+  // expected-error@+1 {{op operand #1 must be statically shaped}}
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %shape) : (tensor<?xf32>, tensor<?xindex>) -> tensor<1x4xf32>
+  func.return %0 : tensor<1x4xf32>
+}
+
+// -----
+
 func.func @cbrt(%arg: tensor<2x4xf32>) -> tensor<2x4xf32> {
   %0 = "stablehlo.cbrt"(%arg) : (tensor<2x4xf32>) -> tensor<2x4xf32>
   func.return %0 : tensor<2x4xf32>
@@ -5700,6 +5716,7 @@ func.func @dynamic_iota_dynamic() -> tensor<?xf32> {
 // -----
 
 func.func @dynamic_iota_unranked() -> tensor<*xf32> {
+  // expected-error@+2 {{op result #0 must be ranked tensor}}
   %0 = stablehlo.constant dense<[4]> : tensor<1xi64>
   %1 = stablehlo.dynamic_iota %0, dim = 0 : (tensor<1xi64>) -> tensor<*xf32>
   func.return %1 : tensor<*xf32>
@@ -5707,10 +5724,10 @@ func.func @dynamic_iota_unranked() -> tensor<*xf32> {
 
 // -----
 
-func.func @dynamic_iota_unranked_large() -> tensor<*xf32> {
-  %0 = stablehlo.constant dense<[4]> : tensor<1xi64>
-  %1 = stablehlo.dynamic_iota %0, dim = 3 : (tensor<1xi64>) -> tensor<*xf32>
-  func.return %1 : tensor<*xf32>
+func.func @dynamic_iota_dynamic_output_shape(%arg: tensor<?xi64>) -> tensor<*xf32> {
+  // expected-error@+1 {{op operand #0 must be statically shaped}}
+  %0 = stablehlo.dynamic_iota %arg, dim = 0 : (tensor<?xi64>) -> tensor<?x?xf32>
+  func.return %0 : tensor<?x?xf32>
 }
 
 // -----
diff --git a/stablehlo/tests/stablehlo_refine_shapes.mlir b/stablehlo/tests/stablehlo_refine_shapes.mlir
index c219232a3f5..c3deb36b016 100644
--- a/stablehlo/tests/stablehlo_refine_shapes.mlir
+++ b/stablehlo/tests/stablehlo_refine_shapes.mlir
@@ -554,11 +554,11 @@ func.func @refine_dot_general(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5xf32>
 // -----
 
 // CHECK-LABEL: @refine_dynamic_broadcast_in_dim
-func.func @refine_dynamic_broadcast_in_dim(%arg0: tensor<4xf32>) -> tensor<*xf32> {
+func.func @refine_dynamic_broadcast_in_dim(%arg0: tensor<4xf32>) -> tensor<?x?xf32> {
   // CHECK: stablehlo.dynamic_broadcast_in_dim{{.*}} -> tensor<3x4xf32>
   %0 = stablehlo.constant dense<[3, 4]> : tensor<2xi64>
-  %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [1] : (tensor<4xf32>, tensor<2xi64>) -> tensor<*xf32>
-  func.return %1 : tensor<*xf32>
+  %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [1] : (tensor<4xf32>, tensor<2xi64>) -> tensor<?x?xf32>
+  func.return %1 : tensor<?x?xf32>
 }
 
 // -----
@@ -598,11 +598,11 @@ func.func @refine_dynamic_gather(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2
 // -----
 
 // CHECK-LABEL: @refine_dynamic_iota
-func.func @refine_dynamic_iota() -> tensor<*xf32> {
+func.func @refine_dynamic_iota() -> tensor<?xf32> {
   // CHECK: stablehlo.dynamic_iota{{.*}} -> tensor<4xf32>
   %0 = stablehlo.constant dense<[4]> : tensor<1xi64>
-  %1 = stablehlo.dynamic_iota %0, dim = 0 : (tensor<1xi64>) -> tensor<*xf32>
-  func.return %1 : tensor<*xf32>
+  %1 = stablehlo.dynamic_iota %0, dim = 0 : (tensor<1xi64>) -> tensor<?xf32>
+  func.return %1 : tensor<?xf32>
 }
 
 // -----
@@ -621,11 +621,11 @@ func.func @refine_dynamic_pad(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tenso
 // -----
 
 // CHECK-LABEL: @refine_dynamic_reshape
-func.func @refine_dynamic_reshape(%arg0: tensor<4xf32>) -> tensor<*xf32> {
+func.func @refine_dynamic_reshape(%arg0: tensor<4xf32>) -> tensor<?x?xf32> {
   // CHECK: stablehlo.dynamic_reshape{{.*}} -> tensor<1x4xf32>
   %0 = stablehlo.constant dense<[1, 4]> : tensor<2xi64>
-  %1 = stablehlo.dynamic_reshape %arg0, %0 : (tensor<4xf32>, tensor<2xi64>) -> tensor<*xf32>
-  func.return %1 : tensor<*xf32>
+  %1 = stablehlo.dynamic_reshape %arg0, %0 : (tensor<4xf32>, tensor<2xi64>) -> tensor<?x?xf32>
+  func.return %1 : tensor<?x?xf32>
 }
 
 // -----
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
index 2f562448e8b..c09e86c1a6f 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
@@ -1260,9 +1260,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir.bc
index 216cd266ae4c0d28a90a02498fdba37bde8c9327..db151f2c95e01506fb5bab84d538b4d9cc0a14e7 100644
GIT binary patch
delta 1013
zcmYjO-)mEM6!-Mw*<E@G-64%7xeaNO)-=<$xd}C-A;wErx-7=fg%E0Ig+g0vskMlf
zI*QI8RP;f{_i1B3%;AOzBF)2IDuViV4EEBuJq$)4WNiE9-ibMIzULf1pYu8A`~4m~
zqJu~D;xz^>%Sr|UfgbJe&C9dV9@9ME*QW`secTN8tfeSurRdA96r8m{N~Oa=bw3nS
zzlP3-hEnhaQ()~wYbhN1_kSBw(eR{N3hyl~f`t8k_yy}&7f3tXz!gCj*>)+yV>TWN
zU){u(aPl0IVvN~eV_S^z+zosqJf6dCfyrcwc9;Tf!pnqj#mL{?L0;Xn?)clCj0g&1
zprAtiMOs$P{=pjs;U9m(69G|sU`Que(GkZ4MUahOj6<H?y6OCbcbY>e3YSVK1%(|4
zj_~9pCPTuJ2qMB$Q<#bhPfueyCR{F~oDh!1po>Zc6}?8aimG3@RzvNx9^si8%;bdY
zb=331jRqPl+-#z02+z)9R@ky&3D3=84o?LE1O%&C6)uu~K&-S*M*ircMRj~2s?XB{
zmuX$S95|cUAjhDeO<d6{A0~i>x<59+CbCpA=&_D!BDQT_rA^8k^u38K^*UzLIrUd;
z0^^iZc6`*03-LYuUeQFxpQpRDLy11RZX(~BPZY@Kj#<)UT8{^t<a4#pXzf<l(Or5k
zexfb@D%lqQGwF$s4u7Sq2gC6r1=ic?4Y|Fc&w$S~Kn3+?_*r{;Ya?|Edi?J3eSN*}
zMixFS@msh4-uXnl_33DqbjSSI-?XCU(#NQ%cG3}jK1p9om#9R>M}$(xm=oOw43iyl
zT*&EwCRExEt3|u0X6#=2Lj7QeKRM&k%T#vhW%au~QmIhIttuVOXFXz-sxDORfaX1@
zMl~0zbwCU1WG0+j@(6XRyM%fN;i3mMsNq754(O74m<g-5nYen9=~eG^+}AOkjfPf9
Y^sQGv{QiAUrKm1sr)Z`1Jo~xz9}AMUdH?_b

delta 976
zcmYjOU29WW6m@sa_D*^UWfL<cxrR1LYns}xo6yEI#29LcF~$-@2{n_>P^PW5bhH!|
z74e&jK8TzrMITgDtRP}uof&Za9mDwIn}{Gji0GW$jF=5)pS@Y@ti8@TAMen|JM`=&
zdg7L4MLZr)hhE>psk=H%^Ow#}4SVWuaL(&qj*{1k(t)-pvde1Cx5LDyum4A{+VK7A
z@<zdW^)9Qhf1=IRtA6wUxex~lyWMaLR<I(FMs)<Mf>De%S`i+$ab0-r7)}UBP9h@4
zuzeb5#7HF1;=FJoiHic0$rQa{3Umn16K;r+y1b5*x?){&$Ms1@0BJFhR=%xinpc&r
zJ!jLxckknYfT%Ihrwc6V9oq#Nkc~&ire<6?m;HiInnTD4XS2w9g(oI3Av`&WNuO{a
zfPio=hg_F%K977*xKKbLEF26%H^m}~+Q!rrrrg4%5=y%}gv(`=6T;Kem`(~;DyXpV
z%nW7>;c69CVatLgT&tmm#{v&Lf+Z{oXGlLF7VA3$Z`&xY2D-a+-q(GUR@5Ke2g64w
zYtTaztJ?Ap1h7#1j|SMpC}j+~r*}0TJY%lWF-jVA%ft!wGHBCD^&vQp1SORn8gllz
z&<p)uUE!EJMVD!v!ku)=M5;a;PLazQv!KV+@AXv3<!JBM+V!?=&*_`^SVR1;NJIR6
zq$xhs_p`QptuORbnzcLakQ)wN23)2AN~^#7o;0S{SEIY2$3Gu<qxm-n=f0cxr(J*V
zY&cZ!ABvMcF+2R37FBg*J7v`RNI>TYBgfU4y@ke9!S>U*x^4T_yq(VEC`ZO-Lat@>
z-gX0q$$82<P`(A)--HTOaG-*EZ4Y{jRCKyx3-O@Z8S^J*oAeTu973suFxQ02RCb_p
z3$&oF$NY3yZRq^F`ZwlRuVNurm8uT2s@}y0T{WsXT}}1I`+Z9!`VLFq|MzNBJ);iA
MCuy<%G``3B57J+>E&u=k

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
index 98607566fe8..95e51eff537 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
@@ -1260,9 +1260,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir.bc
index 60f0f8212995f6dd462de208203b85e60b57c30b..c3bd3005d1c309ef10717b55098a71c042f7961e 100644
GIT binary patch
delta 1051
zcmYjOUuaup6yMp;<7ajW^(@Uo?rrFjtR|Paw8^DpF=R0dA(l`=tX=53Qki3A3ter;
zTILwtALrbQARpUc4|^G#jG=7H>p))=p)cZ-FFuK)0}*``QNP?cmjmBD-{JQ==l46`
z-MUR%x9Q=R7;_va6$*uh%=th1g*Q56S&NZ~VRWD3W;nRwQrL0n&~_J<75!0kzy3D*
zPISygAnbz88C!9&?VtZ*OdqhT`ilKza|t9I3?e94#i~Ghx6?Qy@Zj}I5gyOtW8san
zI47K1L`u|n_B=L4P3AA+vhZXcR|FQ5CFZ~qm=Hce(4q=guc08yy-siqOhykq4l{#A
zf8lh3y^)L<ilR`|`$tN2La&eP*(?g*xr<K(WO@Wqb3w;s>=l$iS=?9K`f=$(<!}6H
zIE0dLxr}mHxKcqyxLQRuDjbU;CS0qbHY!}NqaGKYn!;34I39;78Vxjzjp=Dj2ZftW
zG#f+0trl8&;h7oC7&o)Cm@Npm+i0`!+#KeV@ccaHg&hZu@WKKX@IVkkNU)4$;gY{_
zr+X~+>-P8(MHGEyffRx90zt9RA4(x$DKfF7*LJ^4tNMrChm)tNqUf=OGiE2lG;c*a
zUs7NT9+ee+X0m=ee%{)kvs6%Y&%!zVb399n+D^QNNy_VXBID-|6Hm;8){?nkfv(au
zN=E1d3k99YB)a#KQ407{YeuU3^`15b{P~B?{BHV{$Mj6#U{B!3R8Jt1?hCw?K4Pjz
z>BO!g8;8av_gn`2eLzM1OZq|2LHDNnI*f(yUj50e`1!u0&y76YfkHQvnI*%1fBY}%
z=!+A3siZ%hh?(!!#Co_)Wl}E^$^#MybTMncd~k<?vbaiBpHS5ovTx{(?C{8&eV|Tt
zAE*y>9@5`uZM`#>2sfzVqm2RdsLtf<{PBK!lbSxEIY2nshg#I~q1FI&N^j<DTGO|3
z_Q-ml)~2>kYY))g)8FK5&0a$Po*ULhlT+TXUiC(!%Oqya<>$ZqMt@34-}h?N>Hg*&
GasCCZXub*n

delta 1041
zcmYjOU1*bM6!!G-_%42gddy--z6MQV*JNs%<fGIWV)nI!8cGOV7elQfLqw^e>sqm*
zBK800Mi?(D?82KdhlsM7+cCXy40^N6-OI-27=v97+4km(&K!7u4$pJW^PKlxzenr$
z=<DYgc5K@o2?PRzX8j*tvVG99TA`4^1D%(+8}u)xC}^kX;6MuQqP`X0YhlB`uVYw0
z4<GdfQ(&`(7xhGBa=<sNFGs#xb3np=Km39vED5B0TgFMj7{<D(2#;p*neeI8I3qlA
z4kMyQv*)oUYCL}tmxagkxGJ!iEHMj~z=ZH|f)-V{eiH>zsoR8<fywBh%VB1)=x27@
z@0ga1D2k#`)O&}DbX>0v?OZDgfAJ;05|F78gv|vD=EN>R36#YHwV@Z6F1Y{TZ^I#!
zgk2YIP<V0@lfvaP%3<MX6j9-+DNJn_u2fKo2~SUBIw2g3!6elxs>VjGhMHe^W(G6U
zgTk}3m>n0c*HJGBHyUW<h3Dol$HMdTm{-EhCYr*w4O_U?LJJQC0R#kXw1rFFzzdzD
z(LV-a4uuqbZ-Eqnv5}xy=r5(<vlN+d^vaG;X-U7>aVW7&uA(OvPMVnv)4UaJeMf;U
zj8RF^V{__{W9O|?bei&t?prvce~o47oQ}lzVVv^15l?&Pjrh0bLCcApzd+aNCM802
z%|bz^)A7#zM3{V@)QXYn{IIh@KCk~F)4!8^?+LvTnCJ@pJkk{iC3^xNB@Y|R-z4MP
zifkMjm)v#f^X35+^>4|CT?d`JsSjW*{IvTQ!+*Nx$Xg>XH>1!=r|V?cS4Lmag1(s9
zMJ4?x6E*vKW;N)NOKKCr?UUHAm$Q*A6Pp#3#bqjcgtESnJ)mD?-yQm}2UMuy0hPYW
zgF2my=&iYUrb<;0t@fcudWAKrc~GqnI;Pih5jvqi&qanl>Cx&`_h|J#+KT?s&{p;D
zW<RTAPDHa4_cf{MF`GK$?D4gz<z+2>%y~cDCNa~s-~RHQUiXrI;FM{h^QUvz{tpr7
By(Rzv

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
index aabb800d9df..04227a70c55 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
@@ -1260,9 +1260,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir.bc
index ea1add250b63d2307b2619a00d431ad40d206285..3197c5c29b3e4ca12db870acf3976d422fe26110 100644
GIT binary patch
delta 1033
zcmYjO-D^{K7|!d{*R$p<E3a9o>FGh!G=`&Snx3R>2w9evl7-Th8q!cpDY&hop_W#x
z_%Z4xoyf{Sf9!*S3YJaAU<}WV7hZTHUKafe1h4$K@zx97FCOFa!u!SZp7(j)-`U=x
z?JfHE0TPyFCBxxxm-Y|lH}0M;(_HB8))U@y{2B=@I~1`Tn%Uu?w5)!Lj;P1cuc8SD
zoe>ADJ+bT{@b<kW)Mq_oYNO}g`Z17jD1?w;1uFt+R;O@U;KFUTBAiO&s_>b!I4_)B
zKvIlU`dh4vksbaHmxQy!*btaZrf7#L&?bDGpv1^szJi<>&d&r#!=(ApoMC3LsQXqU
z)LhB1krxAbHPTz4<7%z9e?2d}^#^VXNPi=U>H-a&aX?T2rSX>$@a_AHr8ju3X9xx1
zViCoNaH)in@Yoo}qQbTfTX=jN<9me5Wt3yWl?p0x;aCj1nV7(YmN7Yr$&m2W6s9J+
zgr}!5JuF<UqN>%@YN+Le>vhywcy<=EhVa}R=7cQ^mhk*M=JB^6jIdw{OTvYB2RFR0
zZLRp2y0y1Y#|L{a(25$3EhbM<pk&Y;6Q{NG4uXCX+VzpqaR6+>r9jc3e{`X?V?Ta;
zhR#wTXV9M}&a1_~AzDz^`#wjOhSih4%)8tjf27~;S$yP^99^a>6o_}z4<>S6s(*ok
zKG~`!^L|U5qoBY4h~Dq@CZEy=g3%^{=`;yO9iL#s`9^myoPE3VtW|2I+*I0Wbhzj=
z4NzV=`|me3crW+ggch-R=&inVGj*~}%JU#{D%GFx4i40a^4_I%4^czi7#yI2+8&JQ
z_-b$?SfnBajGcs{y6AqY*4zk<s!i8c58ZKf$Bm>92aS${Fib8}*`F-8shVk_6{`4X
zMIFuzsxLATI;t*YY+6*m>v%%l&De7%Thmpl`qS06U8^m$Ml~O;wV`V*v`%#&t+%1)
u)Iio&jciiQWFzXf&ewE)n(gmhqJZf0mj3_ezqB?LRD9@5YIv(d^VWY(gur_M

delta 1032
zcmYjOO=weD6wdMMnP>7QQ_h4;ZC@XnrZGIlN!ypS4I#v&giuPYF~m|!2#8bZgj%eq
z_(x}^bi`2xhC6?Y`2mYFh%gMkjf-J$({*v@QgG$Z#_cTBo4%>$z&ZCEzH`p^-Fvq-
zXlsN1-bBK+EGrp_M7p)Vv+ZJUw`n$edNdyKp5v!zc*&ut<<Rk64vI_a+um^#fj9qK
zLT&b*3Pv5U_QaAZ^-b&wCe&KrzpE}tI2=Y;u#9Daw6n7~C&(bvE=4#sjPHccU&JNh
z<P{{vNDW`Zsu*M0FR>;(mc_cjWHLoNOo2Ax(*z|(?)nYn#BhEfI2tC+hjtG$gGD{G
zT47hOWZ1}yfxH^;&(mpjp?`2SFT8OFe+o#?2zqsbMSbF+pa2@iT_fPzS6>ug<3HU)
zC<qseC`N@RCNLpfE~DHlY}>GfD-~4s3QtaAGA2AVg{io3ECyXnPh(okn3=&$Sh!k6
zb*fvqRzqz}xL!v+Cp<fg*{tx~9OhVfejf9NaHD~Suw}s#ZZ^@xJwXH!K?^P6!uEw1
zy;F7`<f_~I26Wuqca@e^DYlS2OM#+6_f4GB(z^)yN$BI<M%O{Gi3|k_2K}WIwH3Sh
z;d#19fviD)n7E`C21e<M`gPzi#we?v45YWup7<mEcF*GDALZye-Jn3chrTwE^HPIN
z3i)I!n#}t-ahXE?{u6q?*Pnb!?+8k50@G;|lpLR6-T6%Ge&OujlV`0`E9JJ*pwZ<b
zXd0lrat=IfYw%tk_#Ilrtxw<RiCd{N9a5f$kW;C_gm-ADPL%i7hWm&X)o&w1R8U(Z
zF&$rxtgG+csJiI})qOV_DTR!#LoiG(Q<(zB2ZXZv$36N<g(`ko>Ckgb9ZuVHTwP7u
zs*z5r+vzBsR1ebj{KC$fDpmctYRBB?J7|q+K3eNQ&+MRes{3fY16@%=8JjMsBRXDI
z-(+mn$|OS#YWS@U^-E?X)TE|gHr3-ys=q}6(Kl(m|MMSrI0`C0dW06emC=Uv7HPY`
AZU6uP

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
index 8732686f622..0b43304dca0 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
@@ -1260,9 +1260,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir.bc
index 16296841579855ec0e6d73870637be608437e8c2..faf8586da5e7353510cf8ad10d8557d192ccd290 100644
GIT binary patch
delta 1033
zcmYjO-D^{K7|!d{*R$p<E3a9o>FGh!G=`&Snx3R>2w9evl7-Th8q!cpDY&hop_W#x
z_%Z4xoyf{Sf9!*S3YJaAU<}WV7hZTHUKafe1h4$K@zx97FCOFa!u!SZp7(j)-`U=x
z?JfHE0TPyFCBxxxm-Y|lH}0M;(_HB8))U@y{2B=@I~1`Tn%Uu?w5)!Lj;P1cuc8SD
zoe>ADJ+bT{@b<kW)Mq_oYNO}g`Z17jD1?w;1uFt+R;O@U;KFUTBAiO&s_>b!I4_)B
zKvIlU`dh4vksbaHmxQy!*btaZrf7#L&?bDGpv1^szJi<>&d&r#!=(ApoMC3LsQXqU
z)LhB1krxAbHPTz4<7%z9e?2d}^#^VXNPi=U>H-a&aX?T2rSX>$@a_AHr8ju3X9xx1
zViCoNaH)in@Yoo}qQbTfTX=jN<9me5Wt3yWl?p0x;aCj1nV7(YmN7Yr$&m2W6s9J+
zgr}!5JuF<UqN>%@YN+Le>vhywcy<=EhVa}R=7cQ^mhk*M=JB^6jIdw{OTvYB2RFR0
zZLRp2y0y1Y#|L{a(25$3EhbM<pk&Y;6Q{NG4uXCX+VzpqaR6+>r9jc3e{`X?V?Ta;
zhR#wTXV9M}&a1_~AzDz^`#wjOhSih4%)8tjf27~;S$yP^99^a>6o_}z4<>S6s(*ok
zKG~`!^L|U5qoBY4h~Dq@CZEy=g3%^{=`;yO9iL#s`9^myoPE3VtW|2I+*I0Wbhzj=
z4NzV=`|me3crW+ggch-R=&inVGj*~}%JU#{D%GFx4i40a^4_I%4^czi7#yI2+8&JQ
z_-b$?SfnBajGcs{y6AqY*4zk<s!i8c58ZKf$Bm>92aS${Fib8}*`F-8shVk_6{`4X
zMIFuzsxLATI;t*YY+6*m>v%%l&De7%Thmpl`qS06U8^m$Ml~O;wV`V*v`%#&t+%1)
u)Iio&jciiQWFzXf&ewE)n(gmhqJZf0mj3_ezqB?LRD9@5YIv(d^VWY(gur_M

delta 1032
zcmYjOO=weD6wdMMnP>7QQ_h4;ZC@XnrZGIlN!ypS4I#v&giuPYF~m|!2#8bZgj%eq
z_(x}^bi`2xhC6?Y`2mYFh%gMkjf-J$({*v@QgG$Z#_cTBo4%>$z&ZCEzH`p^-Fvq-
zXlsN1-bBK+EGrp_M7p)Vv+ZJUw`n$edNdyKp5v!zc*&ut<<Rk64vI_a+um^#fj9qK
zLT&b*3Pv5U_QaAZ^-b&wCe&KrzpE}tI2=Y;u#9Daw6n7~C&(bvE=4#sjPHccU&JNh
z<P{{vNDW`Zsu*M0FR>;(mc_cjWHLoNOo2Ax(*z|(?)nYn#BhEfI2tC+hjtG$gGD{G
zT47hOWZ1}yfxH^;&(mpjp?`2SFT8OFe+o#?2zqsbMSbF+pa2@iT_fPzS6>ug<3HU)
zC<qseC`N@RCNLpfE~DHlY}>GfD-~4s3QtaAGA2AVg{io3ECyXnPh(okn3=&$Sh!k6
zb*fvqRzqz}xL!v+Cp<fg*{tx~9OhVfejf9NaHD~Suw}s#ZZ^@xJwXH!K?^P6!uEw1
zy;F7`<f_~I26Wuqca@e^DYlS2OM#+6_f4GB(z^)yN$BI<M%O{Gi3|k_2K}WIwH3Sh
z;d#19fviD)n7E`C21e<M`gPzi#we?v45YWup7<mEcF*GDALZye-Jn3chrTwE^HPIN
z3i)I!n#}t-ahXE?{u6q?*Pnb!?+8k50@G;|lpLR6-T6%Ge&OujlV`0`E9JJ*pwZ<b
zXd0lrat=IfYw%tk_#Ilrtxw<RiCd{N9a5f$kW;C_gm-ADPL%i7hWm&X)o&w1R8U(Z
zF&$rxtgG+csJiI})qOV_DTR!#LoiG(Q<(zB2ZXZv$36N<g(`ko>Ckgb9ZuVHTwP7u
zs*z5r+vzBsR1ebj{KC$fDpmctYRBB?J7|q+K3eNQ&+MRes{3fY16@%=8JjMsBRXDI
z-(+mn$|OS#YWS@U^-E?X)TE|gHr3-ys=q}6(Kl(m|MMSrI0`C0dW06emC=Uv7HPY`
AZU6uP

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
index 4b6ebc574ab..3f61d319ff5 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
@@ -1260,9 +1260,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir.bc
index 6b9a004a808b7357a1d4c628863b7a210b96ba9b..ea07a2c9baffc4dc1acd00bce95634c7db2c34ff 100644
GIT binary patch
delta 958
zcmZvZUufG^6vuP=&GEPNN36#*g!H$VCfU#Sr<+~!OQ|7-n3Ym$C?(de#2TiPO}bi2
zJL>34om00V2r|$9pfI*F<{yUQ%!5z9Dq8j+;)D3)gAXb`)EB`A_5QNW`QU+b&pns#
z;rqGwZeOA8EA+`1a2>}Ph{a+(CjX6Fji2tZ?E{H~QM7->#dvhdqqyVI{w@!-C4DRT
zsD7J#KIwXh#67S{cgaKO@7v*OwQn>mX)mH@lQF&2ckS>CAmL~fQNc2n1rqFzVpZV7
z4_Xl(%Hyo?u{Ep<4;;XNSVQ>@oDgex<Ro4b9v;Cdfz4!#N!S7t!t(?bt90fKl=K7V
z^=L53NTDni%39i4p?N*FbJvNo@YQR0UqBQXCd~#5rs82i1(e5!R_Io|`bzCD{Ao0V
zig2}xYFxNhLrr*e6r)MuR0=8Ku`!ILg~!J+o)Mmyz(h_slYwdKb<_>U<Rm7e!c$Y2
zs`m&t8fc6NPfugo(9FzWrX)N&i&++KHqo?%=jJdc>^N|QTP?J3T@XV|u!u$B3YmAr
zLi<pv7j#pP_ovPJZvQI%sP|=_bAJleEc(>Osv!>(z(V)$u)^kxY#a(zE&AAouNN|>
zx_{dal`MMC#=5#$pFY=S_OY({MZK6U1Z6AxC4HlsE6Q?TSGFsZOVDW>W&*4a6chSv
zCajO|3b+5rc@*imt{B(0=N_a;$LMKe)IRs%7TtE*7dXui22Pg-J5E0gwoFs<a)~k<
zRD;U_Zo~@jLBzIrRgZd~8J70$hgM)<F7Nr>-1Y58j@(81rTc%7b`MpF%Iz2T-b=Kg
z>-iyC)9d*(Rdg%=@{Wy-kaZ8Cs*OUgP84FaPg?~?&lPHV(sWPi&4NSw^=EU|@4LO{
z&v)eY{}3M3ito^2ZTgNJ=?)&(vUzi}qig7<->dKVF<Q}&jOr!*%BYUtQQf%FQ8iVH
pdo@yw=~G2lYo@+g%sJnO#DFc{{kaQw#@6)f;uG{;J3jo3^AA&0%=-WU

delta 985
zcmZvZL1-Lh6vy}Fzn_!Iu#K<dvTP>nW_LR#qiy5vOh}iwtlK4oG=vai8e$sO7Spng
zp)?_uw4^m^6;HCh(n29vDAkrWBJROcuN4eE^w3l3!J~?Y6pG+M{dR3^4nBDEeUJIg
z{C|Ak)-~F?M&I3o>pPCqZQFLc?*FPeB-<^k5{v2gcKtV8iAEMYiaH)0Zt+lBkT2o~
zEQJ1kR<7KOzh_20DBbRYlsn3;rt5)|o+-7sE$baOj{6{1kq9D;Q#i%o#_lwh7=7q#
zjLfPxiyEutGgx8OeGc7h^k!FafsOwB8a`&#pT|XprO0C4uo${nogfGs1D}140lDdX
z67lty3JDb0Ktb}I1v(*TI`>>Cu)2N&-!O<8fpNV-Rj+uBQ3Pdi%Lx5jAAMN*3x8^c
zP-Imqp%i6RE~Ct9a1euWR*3`>tcHd#lw>tLjNufkkr9k!Sfx_X!{{hRwZ_;O#v-i7
z$1y(A&T3);6aB0vCows|YHA8oc~;ZYm{zQ2W-w#0nw`ZgE60Jus!~A(w;48U#ysX(
z6-l2Xs`a_VOQ1P9+m+PMt6j_VSYAvWb$<?(4Eo-}k~R+$pqO^I8ex4k7OJ75LEl>F
zlgp`#ZNK~;${X~xg%xqrK0TG`^cyV;r{zjI*H|{wKhk}XOu)+>E$Nm}CPtrH7%&J*
zUJhdNZYnHq?g`ib$arKnQx|n=-E)tT*$jF|2h}gVvPsX94mOgiJ&mOGo@Uaco>}dl
z_cF19(yCflHFQlQybq>js1;fAe$X!Uy{|4p%Y3!(iRQoDfAYD~uWkRMw70iNRH(oI
z`YxiXoXhsos@%>dsVJM-^D>vS<w(wywVX`{<$lhQOSzJ4<?O^;ZH<F6ELE$8j28&Q
zJFE^%!FQy~cjc&WCys5C>+4N&YzKK<YQ94!Wkd6Kd{Z81esP=M*l6-oJNOTT0*B5@
wK5(QnaAhU1<?X<2y?eKbR%9*M?>r2#7Mg#4sh79^xGFydhv-2)+IPhH2RA9rT>t<8

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
index bb0d581b63e..86f49ef22b6 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
@@ -1268,9 +1268,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir.bc
index 4ec2df8d1a93e56a1ba089fcbdebfb4deed45d9a..4134975b8b1a54cd524b146ba5b7900704b22ada 100644
GIT binary patch
delta 907
zcmYjOU1*b67|#1WIetx(tbOyfHrclr^L+^+6G>~3HfC9tmZsKHYbmwVtcu7OwZu@Q
zlOR=t8`7=2a4+m)yBuY(d7%Ua5m8|<*zjU7VcXSq(TfRQh~m{mPvZPJ@bdnghv#|T
zbM80jeuFN5W+ZLfPFa@KrQ_9k(PU#?et)H}Pa~Q?7=LdG)Y8<Orb9*=rJ8*08<M|#
z$0Erze7$KX9m!f6o@f7`q?9{{WxexO`C|~PU=TsZGL{+Ks!n2skwv!EGONr0F0wjR
z$7xonBS^8C8Cb<xHg^v;aGurfLAmJf^)C{H&D@1G<m8I~(_jm!5JR3#<YhQephY<s
zNJR6jzT3d}45F5RtqUyaj2(=E+zEuw7FZRFDE6`{l~7_eGJ=r^t5^&%R->aB)oE-D
zV{umF;}~~Y#pBS;!~`a^=Hw(MgRG{eFxACsdK%M%tY&5~qm^f8F`HvGH-|aJYJMK`
zCaX#X6;{zGqO2Acuyp}HGAvk(Dypmsq#p@O%@e`LM(nPquS|;|lc3-wnCbv6ncfV1
zev{~~G`ef_k38)D(Ao5q%5(uMdOeMU#;uv&6qFwyd5UHF0$Hi?8Q+74o?Mx}#c6rM
z{d7}Op&?_jDn~<^R{AP*gMO3!)~>+v+lFI!+=KKpa?N$?1L|lCY+|MPd(T1gwMjb-
zWSiNozf$K5+T(`i?FjGEy!r5ZEy8;Eux|bcyM1}3U1?`ZJNw2ky*pbgs1*rCu4!QN
zm&hm3(xvF5_Ewhc-8L6|W#V&a7nimjBg!|sx4%WSB(Zpg>M|9NQ9)|)Gl|tzkGX|V
zd_yMWHN9z=e${BS?@Gc>kN6x*hMbZlbzXOTQg<wRSH5&?`N>Jj6+Qh=Pyf}^J@VAC
z>3xZ~HV?n#+Sq7w&B#%=N9J9N_DR*XX}_$ww%m4;9apcm0Tp@R_RFSgNq@q%|Majn
XSAG3C$6xZ*<wRm1-D_S?47~FkP<q7&

delta 916
zcmYjOU1*zC7|#1WdH40}m#lsB*)=uGV)8YFOf>6CX={uzElq7p*M-uRT5B0PMpm;h
z%9J48EGlDjH}yp?wyQBjP%kV2p~%32FqrH@P}t4DU3lSO7vj&ApeLF8IdFK-bKZyN
zdERsGZPL9>`uamFJs69{GM?vknf++BL88lbmwS6n-)nwh{jw`m%TiC4UbV6))$~(m
z#6{rW|5sXXId3{WS+LpZnof6)Sx#CnbY3qHf`r3igavif1=2d5!n$AxLoF!6xdChm
zpMD1o;mo_ph@BhQz*(_}M>p}l@bIW^xjpW2f{L9#cOH4&c0ULY8YUxwg4ihNSg1h9
z^}SGkydZq_Gkh)}Y8k}L0aa7+te~jZL(#KE;Zg~u9^vtEj0;apV4_brkw8LtauSoK
zOif`bDLgZS8DBV=gaNa&m^GerbC?SYS1PD<3D3`Cepq;60SkHI#YHTR3NI~TiG`P!
zv1|*ktYAer9!FewbrpM7@uk3nCs@OpaFNU-q1rqd{?kev3-p$$3bF}|Ai?Itv}Ol$
za9o?{m~M90>5l&1{aWg7pj4)FSU2Mt9t`ii9n64TR0G8_eS{&c%Nggl`+?CiT}DIy
zmIL&cW+Ee2?UWvk<XYw9$hUM$4|xYer=PY`R=__(H;`{$^j@Zpw!yDhZ~pYu5pvq1
zofd|ggL{6U&Iht5ETh{IJ!o|E(HC2Sjp*wp<?X27TVRvQBx958l4S><YmLASjr4tE
zh|QaQZ^KxZ;(xTK^7NH`i{VEhF2JPtYTp|~g=Y6NyNRlrNakom=aLC3YCU;Ios>rr
zji+L|H|1y}<>^SOq@S32(co><-LH2|_xn^@+4uAxbDQ+N_zOWhn1$`~l`8>z7oqY*
z&X=^}$7oDzeoWu;9o;hg>4*I7?KXeu3H~9yW0>Fh>5iS9c5_AV`v*F%U2B6?P4_<^
c`#B)SW$p1dop>m?p(p#N>Br_y|G?h=0KQMeL;wH)

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
index 8c1f48bbd8f..91038b90978 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
@@ -1279,9 +1279,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir.bc
index 81b4be150f51ab5d1a4de465184bb80ad52f6864..7d52c8e3cc44d22b07ac4f9771916f0692986ecc 100644
GIT binary patch
delta 1021
zcmYjO(Q8v@6wkT8d|uPrtUbAH)AlaCCO4r3*-{&%hAf8IHKmqXyHdNFRUgVWy3&<e
z>L}8cx`|Yg;SR!HWH>2fsGD+edl+MXKvD3)hrJAZ5K-`9>_L3mzQo1lz~TGOJ-_q&
z{myrvouy}I>Gqb9avZ1Ew(S<3uj5sbbz3Z}5{+tv|D*A8SMySu+S7F0NTakQe&~R_
z2+eh+(g?Mup>(E}(g^(XU!-KN>5%-;^!rC2gII;b2s4(k%xJ0JR<O#*B3mt6SY`I%
zGON$laE4XyN%XRr*|&}jHun!~;tN*$2XKyIDY956EQStN^8{ftcYX^wxo>?Lu1+er
z$g_#OoNO-8yj*Qg#`3JLZ{vFgQB}~<7c4Lw>}3?>&*sjJ0;^&X#dcPu5=yKF2Qk>i
z%5~wg8XCfoK8A-e9A`B$f)S5ZJPuuqj$%}6j*Vd~%xZib<9eHk2}}&Inw-R>R-T%|
zRF2j3G^Q1+nHkKOtSS{$SjA$9v6`Jl$1H9zY}kxNEV3$)el#rj7sD?N_g)}crg@M_
zPzVxCb(9v&U<M(}B)TWpTbJmO*!)4_aiCPD^H|mEX&f|e%MA8HS@<weEYn3~<uZR0
zdj3}+SEg@qM*1Ut^qb5@4jA)`ay62v9`_@^(gV3*_qQ&08wn%e9ituOs^apf&GDbu
zaoSxIdWu!Qwf%hx)yN|Tvi@?%4mG`^KV@kCa%aEhU+a9QO8iUb30-t_dC|Pm?zBUt
z9e!t+!MCcbs8u=I^|J>0ySq<ATiA&Gt+!k0`Gj8^`5FyQLhidndH=|sH;5KwHlCq1
zxfFM)Ah+V5_pPr7%w2?HgH(6W45ndPYIC#JD>d|bq?mBzbfP4437g)QpAruB%Z~nL
zc&WB4S84;J4Fd<o^&C1RdC$p}>%`Zt)o>FHxWo0{_I9l|+0Z*8Uull_<c8<;9INB*
z+^OMa8gR#h`st%>`n68fQFrgwP!-vm>=Q3(%cn`t`6Ix3<;DMh|EW5BO>QO+(?kD7
Ia^=l`0R%9}<NyEw

delta 1017
zcmYk3PiS0K6vpSgZ*E^EFJp4%HJK#O>2v-OhSE-(5Yi!z!!%|>8qzd|G)>3pEcz!*
zLkQ{6N{|?=N{K}zo{gIzwMZ8=w&j6Rqzl)@f{06nuA~81BBd+6^K|N5INW>qe)pX3
z+`D(3_O8?YEh9bTIL?4=+nqXIE<bTQEvp=jYTfp}F<$KJSjkX#hRzxp6j#Ix9kUSl
z`@fZzz0ieFcLqvFdPSDnj~k)1+-d*e(h!JMIE*l36|0QS#%T?8#xRB(X(y}fQEaoi
z`VQ7v4ZMp1HnT@Ju*v2~?tOf~Y9xn`7?vW7b--fi#cGKlY>wX8!l*p4J`4}(n+h)S
zY$7k!jyx^NPDd)1XLb84++h$k8an!b8pFXsMnQh-=-n)^Di%@fW;Hg3F;?T_81G}{
zx^P)dOkhH9larW?vznU1RFYLZ4xLO-V_Nr|nZZn$)$A;0^)z#Hm>Xd=KacrQRtpPQ
z$gx^n#G+zVE~9L+s#H*66^kLps#--)6?YjnY{oK{Srte>8fxCn@C(De7>Jgr1~Lf>
zL4v7H(XtuLB4n9F7v*->D*Y-pznyp*D3<63>UuoQL33MXZ~)3eEl?=YO$^I6e;#`F
zM<7?CFR(6ykt6h%oR1tcR@P)El5K2HBHz=G^07VGb+z9}7=h$zdW_LVb9uVo@qV}C
z)aEyOhPv0)eTG7Q@rZ$8Z>{Gswf|fGilOV5dIxp=_1@PTg}>`Pr``7SC8K$z-D!tP
zJN(ixgNGWUsC7Bj_pJuKw*G6-7B*vlYQBEpGQTwP8Csl#+&78x-pSWqA*x9wo}~@>
zB<@l{cH{5KM8c+iDJL8{AomiETuT(?=Y$<Q95jO&n3me!4w(B0lPz8b#Z5Y5B-1jM
zv~zDX$(<daoM|DCH?>>0d~LQxE6Jxxhfc`dq?0?@<aT#`ZlQ&HtEoMB;A@L5+G%l8
zj&U}qpFFx>ztXuT`|zR9R^(vnNZX@Feo~e5sl(1M0oLOz|Nr+_ns+wj>(mMQ$$Orv
GzxofP1ITCq

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
index 152f30bf6ba..e6bdd26b7a6 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
@@ -1296,9 +1296,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir.bc
index c9ca8fc333543474550e4030dd31b4ee313bd7c0..f6b993b17c6be8e5af166066f84452056d827a6b 100644
GIT binary patch
delta 1126
zcmYk3Ur1a>6vp@5Zzq%M&AL6_Y!<zk?Yds&#v1!4LByJ5yR3`dikd}5jkRn4k(4$x
zVp=1{nqp-eY+?pNCA5K12_=M97AqwWA<z`k_OUN5_#m`}6nyAIUkYtHcatR^n0w~V
z_nU9#ocW_bzZK}~&x}ULactWT>+v|Ec2qAu7mG#{VapneYNq&|aZc`@@k}qDHw@`f
z)T56K56KxdTyaQ!Q88Qx)ANL9I{GtiN)ICH!JY9Cc<?AZ^@P={UbB8Ub`r!Q97dQi
zi&+NyvsW?4Xh*yM%PbDI;tGoo=dr+|>2)-*JlI;md6pfA7jc0_$6+inEI}6Q0gIu7
z#k=a`T`J`B0tqBpBB}O;I%X#b*ZniT)Cl1Qo0F47;^tROa9w}z1k2_(k7Jc({5?Ya
zVJfnSn7*xP!t|*mi?wxp%OLXaS*s_Q(its`6zCb;l7WYD?t|oi_(wB@6pPt(8tDj&
zOa>ViU0vv^W|2rB!6KVQHpZg68{KghJw52ri>?coMPDEK8d>CW$c0$+_oF||VsH?H
ztt^IyFx1XsWCSB^EJjB$Dp-t-VN9|Jg`oLZ4B9uJM;>i9BZ3HH8q+LNr2CI4^?qfd
z;?HqGy+ktjeQ~AoCnL5Ia60KYND>5s1SxvyEg5WqX-T3Db&*ffzn7|x)!hxGJ81!P
zx;~9TV+9#(gs?CfNOjVAw5zjKpIQ}0l@aLZq^np^aeJR^g0w{cJ(aWDjHwy*iQVF-
z+x9KmEKcvfNjpovTbL_8y=R!rl2?s^_M(dZNW02zXAEt3qWUS>{K)R}>JHy$qxxmK
zr=G7#P*S~86QMbEy5<`#6!$wRgnp+$^8oDtt_y;qJ&s!5dxY+)U-v#sn<^ALVBEj2
zQn8CPp?;2?(#2+;{DrnuJf5Yk;@j~9M9Jc{x=NxcwVG(5d9|5{QA$z$*&~HQK<*%<
zwXxq0q1$mND1)AVy~V|nmf5DgP}Y`~N?Mn?UH_6=ueWQDmX+n@lG43x*MN$-jv8=X
zm2&Oc@iMuxQX+HP$T!pl*P#>YhU>JwQ|4AzOWe>l?o?S@TPtZJ+qBc_f%Z6~oCZf7
oZ5Z8o=S~U9tM!Kc&ffvn<xW3-{L5vh`Ady2(qF}F<CWU~0aS+HcmMzZ

delta 1055
zcmYk3UuYav6vp?=x7VA=bhDgH;*yN%xSdSGw2h0}5@OhvW!Z!nLK0&PX&RSJ&>9QY
zSYu2=y0#L_T8p(8|53%h6!Fhe%TPs1g@Of(#Rs1hDN?XPANrtginiX_8rlPwbMLvo
z^WE>v{8pjgEA+*OM(4wh<Jh*H)ay>9Qo81*QeM)sim8YQzcmiYrj@`9N+rXP0i^<Z
z%?OZNQM>AQScu*FKLYh`{Zyj}OeYY*j41<%Q~-ZPHCZEu8K@cS^6p*`i)0c>#wu1B
z9POS$nUO^{Ix>sx-8jMG;Hy|;vE?wfu-x8V!4Z}_b4T$ei=8>V#jpfftQRbX4i@{=
z>l;*}SMvfM@+^^8n-kgcJmKT$jZZW}_>9f{`-w#LG!uM3>YZoVwPy}zSf&>U>Dwt}
z(QHPtrU^5m@+{7Ng|8Vzk<T{0!J^K{FbbdtaZbi=$DvnpH}SV-2n819{(kf~vltk_
z0E@vv47RZFJa{aIhA`x^7#_xOn#JfSMs?Bm;j<VU!&oPa@o|hNSWHY{BFSQM5|iC5
zrlv5}%VK&O(>*MTMHB^#QVAu=B9VaRT^Cw+b{4bfu^G*1W-MWeMS=A9$D(?<!K?ps
zM$ib6jQ<dxYPf8;3o)mU=0K7l5GP16LVIOA1Jja33+fo}qrZ+f?%s4MR_LR{DC_w&
z28~r@yc5F0QY_y`N03zq8sD=Fquz*R`{*OAskFVht_#u<<E@5z%<eH(52~YfCQ9Gf
z=jckf+*GByn(z&j!_2+A$gGLA8pwvnQ{U5ub-86j>&>-nt;?+IebADPgnnsxQoDSl
z)l21twxz8JZS7M-#=D{s#F~1s^-B$gTb(@Us(OCo9=fD{+V~J%QI@;SxOPqD++#Ga
zeso{ZLWTQ(rt8X057PB;Z+aV1K0LXpfoM^k_A*pazj`he)GhCTdcWPK9je;ysL$I?
zW%{=Iu05xGo$vEamGf<SLe2Y*dffNb+q(LOuAcGjwx{DV9z;V_t1)>GVd5^^&#J4w
zL(i#yb>BalPj}c|FRk<E&)4|kUHtO;-GvLayOJt&JX3e^V(l=i&UQTD{1IdQ$kLs^
U_Re}wMcwRplx~CrohRD<1AQ3W(*OVf

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
index ca3ed171335..05864477016 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
@@ -1260,9 +1260,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }
 
diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir.bc b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir.bc
index 8b33745e37e76c87e1d506ed78c318592ef8d662..b461ab54a56e0d727552a535636f67581433c329 100644
GIT binary patch
delta 936
zcmYjOUr1A76z}x&>oRXk4s)#Sy5=@@-purFn9MQPWen40L&sW05IWW{6Sc625Vd+L
z>g9VaBr1A{AR0Z{bML+eQ9<?QgRfHAx2?zn=O5>H&hLEZ+k8x$kLk?@LWW^PTrO9O
z(uOpmSfEAMFL^wwT>XSecgu`PZo{Op{U+=gmnpZsZkhK6<*K*n?KI(Vn-EGnXH4wd
z{VSdFi#OfG0n?%2E;;Obad`^F!s&#QF^gFS*P=PhGh&F<QfAR-;Q@=q>sVqDIgbcy
zeb!CfVr^jXHkMfo3}S_$3({FBbcPZZCBj42;%j#im!rl#XRTAfj|6K-$d5*nO7d`9
z$4Y|5lXX005Y-61DnMCz^e~bj3(vHD4Y{*w|Hh8O5RxoXDWu#iY#TO<bQ)<d3%?(J
z7MTn(0Tv@87zwf%9mQyvMKB1JWV6VsHO9v=?qre6A$OvM#l!?A1dDtgd5uM(fC3A{
zfWe|zL=i6-F1Q#KR9GZQoebsbcH7o|nv$Ob0cCds3p6Vi+AoCX$kynQj(N3Y69Iy0
zN7Ep5#3-fFedXm=`?9`B*D0>i9UV(@DQMAo`7$_!0UDIwg2(IL9eS^qh$!Nr%Q_y)
z`C~!(E2KH&v_^L+?4fBL@hXQC<fxOLQKZ#EN0A(LrW1;(`mXaW{i8Zkqq-KUQ9X||
zsJ=!ns%*GBL<u?7-Ky-3?pHPPYKM6e>VL||eyG=-?w$V6`k`6<Cx=7TPrWfx9o+Ym
z%JNWu4<+SLzhBvN{Wnr6N|APekP_+!>JCEPp{rSgrVE?wddc2XI@&0uDP1q6rPDek
zo2*tEk~zyq!}5ydlj~N<ouzCYk=@fgCO0f!+nL6Q9OdeO+#bNB+*U)*%7am#w4)(M
pfeLj@L0*ma%2KpdzK#aG6=Jn;<^R83YfMba@6imEt39zU;}1DYvZ(+7

delta 915
zcmYk3Ur5tY6vsP$yZu<dmK@fw?dO`?oOAw3e;aDL*7_kL)27&1b5VrK5-IT?8WCb@
zPer}l*FyS754}XAC!@#SeGMc6>&*vWrL4Q(iac=6y@&I?_j5RB>nUwLrFWYM#tp-W
zxLmGk<rTAU{%T#nQB$LeuF^IxyPcC!avM<^sEWdxl;=I?bR7Bp*TQnaGwN_hA(V$F
z<&LMf$`O`f@28t_5DTXhPR10b7+k*H!Zf1;9p#i+w3&FsVrC8tEFz1Du-0bY!Clt6
zy6<6`MOQag7`h;xl|yGJV=+c}%vxe~9SNB>9ysIbrGO7f){vB6j3kZ8)3x<0NfysG
z@SH(Z271&56Y4}WBLy<?Qae(ywfoiq_LYW^VqsaZ+$?%~(aWN*4}BgMJ|BE6`uoxE
zXOT`L9bhpyh`|txKmaPqWROv93=d=2$zo&#Bc0VOvRPyWi(C#ljYU3>JPX5s!J<$=
z0k0V@xEMtgS)@pf3=^fj+TALO%ddXFiogA{G$m*2u7+-rrO^`|)2hj%1PG>mO@q+U
zK`D(Ms#EUPE$cHhN8K8&=~$490h1Qxn?Mh`s9Wv?&f2Ft_)#AtqDT$h)bUu(v;^d@
zpyo`_Dy>tfhQ@UyN*qd%!-ftiXen2pCx=byRHV{}@O%1)I#Wi?N6M&IkqYWt<eF;!
zLSvAUa;))`ic5`e%kWZt^di)LiY-5s;8N@Of9mJsYM&emmbP2tq+W2_PnwV?+nXsR
zd)j>}j<zqkEwV^EMzDmMftrI*b7+tU?Otk<!=_iBHj{GJJXPCsT!W^IG^K5M`jGsB
z+*HC1Gw9Ax#;#@#RR`pr>7^lgBId2VT=^tRSzDGpB)cN7$Gp;t1s!?H+ts|h9cy(I
is9<LW`6AZjDH5w*i~m<PU->a5zsLG$qSPF3F#Z5Lvac-w

diff --git a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
index 7840b5013f7..c19d4810fa5 100644
--- a/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
+++ b/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
@@ -1301,9 +1301,9 @@ func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tenso
 }
 
 // CHECK-LABEL: "op_dynamic_reshape"
-func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
-  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<?x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<?xindex>) -> tensor<?x?xf32>
+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
   func.return %0 : tensor<?x?xf32>
 }