type: DT_INT32
    }
  }
}
node {
  name  : "y_reshape"
  op    : "Reshape"
  input : "y_const"
  input : "y_shape"
  attr { key: "T" value { type: DT_INT32 } }
  # Attribute TShape not specified; needs to be set to its default
  # by tfcompile.
}
node {
  name  : "x_y_sum"
  op    : "Add"
  input : "x_const"
  input : "y_reshape"
  attr {
    key  : "T"
    value {
      type: DT_INT32

# Text form of tensorflow.tf2xla.Config proto.
feed {
  id { node_name: "x_const" }
  shape {
    dim { size: 1 }
  }
}
feed {
  id { node_name: "y_reshape" }
  shape {
    dim { size: 1 }
  }
}
fetch {
  id { node_name: "x_y_sum" }

// 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>

^bb0(%arg0: tensor<3x2xi32>):
// CHECK: {
// CHECK-NEXT:   version: 3,
// CHECK-NEXT:   operator_codes: [ {
// CHECK-NEXT:     deprecated_builtin_code: 22,
// CHECK-NEXT:     version: 1,
// CHECK-NEXT:     builtin_code: RESHAPE
// CHECK-NEXT:   } ],
// CHECK-NEXT:   subgraphs: [ {
// CHECK-NEXT:     tensors: [ {
// CHECK-NEXT:       shape: [ 3, 2 ],
// CHECK-NEXT:       type: INT32,
// CHECK-NEXT:       buffer: 1,

  %0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.reshape", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
  }
}

// CHECK:       module {
// CHECK-NEXT:    func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT:    %0 = stablehlo.reshape %arg0 : (tensor<2xi32>) -> tensor<2xi32>
// CHECK-NEXT:    return %0 : tensor<2xi32>

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

  // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%[[BROADCAST]], %[[CST1]]) : (tensor<1x8x6x1x6x1x1x18xbf16>, tensor<4xi64>) -> tensor<8x6x6x18xbf16>
  // CHECK: return %[[BROADCAST]], %[[RESHAPE]] : tensor<1x8x6x1x6x1x1x18xbf16>, tensor<8x6x6x18xbf16>
}

// CHECK-LABEL: avoidSimplifyBroadcastReshapeUnmatchedDims
// The reshape splits broadcasted dimensions, instead of eliminating size-1 dimensions.

// 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

  auto y = Einsum(scope_, {x}, "ij->ji");
  TensorShape y_shape({3, 2});
  RunTest({x}, {x_shape}, {y}, {y_shape});
}

TEST_F(LinalgGradTest, Einsum_TransposeBroadcast) {
  TensorShape x_shape({3, 2, 3});
  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
  auto y = Einsum(scope_, {x}, "...ij->...ji");
  TensorShape y_shape({3, 3, 2});
  RunTest({x}, {x_shape}, {y}, {y_shape});
}

TEST_F(LinalgGradTest, Einsum_MatMul) {

namespace mlir {
namespace odml {
// Converts mhlo.dot_general to tfl.BatchMatMul. Reshape and Transpose ops will
// be inserted to convert to well-formed matrix multiply; i.e., mhlo.dot_general
// -> tfl.batch_matmul(mhlo.transpose(mhlo.reshape(operand)), ...).
// Note:
// 1) Reshape/transpose are inserted because tfl.BatchMatMul requires
// size(contracting_dimensions) = 1 and size(output_dim) = 1, whereas