// Confirm we can extract type info from reshape

func.func @main() -> tensor<2x2xf32> {
  // CHECK: %[[cst:.*]] = "tfl.pseudo_const"() <{value = dense<2> : tensor<2xi32>}> : () -> tensor<2xi32>
  // CHECK: %{{.*}} = "tfl.reshape"(%{{.*}}, %[[cst]]) : (tensor<4xf32>, tensor<2xi32>) -> tensor<2x2xf32>
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>

  %reshape = "tf.Reshape"(%dq_pool, %cst) {device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
  %q_reshape = "quantfork.qcast"(%reshape) {volatile} : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.023529411764705882:-128>>
  %dq_reshape = "quantfork.dcast"(%q_reshape) : (tensor<*x!quant.uniform<i8:f32, 0.023529411764705882:-128>>) -> tensor<*xf32>

    %cst_63 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

        %cst_63 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

      ]
      y_signature = [
          None if c not in contracting_dims else y_shape[cidx]
          for cidx, c in enumerate(y_labels)
      ]
    return x_shape, y_shape, bias_shape, x_signature, y_signature

  def _create_einsum_model(
      self,
      saved_model_path: str,
      equation: str,
      y_shape: Sequence[int],
      x_signature: Sequence[Optional[int]],

// due to the outside compilation attribute could be removed in
// canonicalization of Reshape ops.

// Reshape should not be executed on TPU as all are marked by outside
// compilation. And there should be no host-device communication.
// CHECK: tf._TPUCompileMlir
// CHECK-NOT: tf.Reshape
// CHECK-NOT: tf._XlaHostComputeMlir

    %0 = stablehlo.reshape %arg0 : (tensor<10x1x3xf32>) -> tensor<3x10xf32>
    return %0 : tensor<3x10xf32>
  }
  // CHECK: %[[RESHAPE:.*]] = stablehlo.reshape
  // CHECK-NEXT: return %[[RESHAPE]]

  // CHECK: @main_1
  func.func private @main_1(%arg0: tensor<3x10xf32>) -> tensor<6x5xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<3x10xf32>) -> tensor<6x5xf32>
    return %0 : tensor<6x5xf32>

// CHECK:           %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK:           %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK:           %[[VAL_8:.*]] = "tfl.reshape"(%[[VAL_2]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>

     // `input`. In other words, the shape of the `Reshape` op are not
     // changed after the transformation.
     (IsTailOfShape $rhs, $input),
     (HasRankAtMost<4> $input),
     (HasRankAtMost<4> $lhs),
     (HasRankAtMost<4> $rhs),
     (SameElementType $input, $rhs)]>;

  // Move binary op before reshape:
  // binary(reshape(lhs), reshape(rhs)) => reshape(binary(lhs, rhs))

// CHECK-NEXT: %[[reshape:.*]] = "tfl.reshape"(%[[conv]], %[[shape]]) : (tensor<1x112x112x32x!quant.uniform<i8:f32, 0.023528476789885875>>, tensor<2xi32>) -> tensor<1x401408x!quant.uniform<i8:f32, 0.023528476789885875>>