return %arg2 : tensor<?xi32>
  }

  // CHECK-LABEL: func @main
  // CHECK-SAME:    %[[ARG0:.*]]: tensor<10xi32>, %[[ARG1:.*]]: tensor<10xi32>
  func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> {
    // CHECK:      %[[RESULT:.*]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]])
    // CHECK-SAME:   Sout = [#tf_type.shape<?>]
    // CHECK-SAME:   dim_args_spec = []
    // CHECK-NOT:    _entry_function

  %arg0: tensor<1x64x28x28xf32>,
  %arg1: tensor<64xf32>
) -> tensor<1x64x28x28xf32> {

  // CHECK: %[[ARG_PERM:.*]] = "tf.Const"()
  // CHECK-SAME: <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]])

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[ARG_TRANSPOSE]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NHWC"

  func.func @main(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
    %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
    func.return %0 : tensor<f32>
  }
}

// CHECK-LABEL: HloModule main
// CHECK:       ENTRY %main.{{[0-9]+}} ([[ARG0:.*]]: f32[], [[ARG1:.*]]: f32[]) -> (f32[]) {
// CHECK-NEXT:    %[[ARG0]] = f32[] parameter(0)
// CHECK-NEXT:    %[[ARG1]] = f32[] parameter(1)

// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>, [[ARG1:%.*]]: tensor<f32>, [[ARG2:%.*]]: tensor<*x!tf_type.resource<tensor<i32>>>)
func.func @write_only_resource(%arg0: tensor<i32>, %arg1: tensor<f32>, %arg2: tensor<*x!tf_type.resource<tensor<i32>>>) {
  // CHECK-NEXT: [[READ:%.*]] = "tf.ReadVariableOp"([[ARG2]])
  // CHECK-NEXT: [[CLUSTER:%.*]]:2 = "tf_device.cluster_func"([[ARG0]], [[ARG1]], [[READ]])

func.func @testValidIfRegionOpWithMultipleResults(%arg0: tensor<i1>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0, %1, %2 = "tf.IfRegion"(%arg0) ({
     %t0 = "tf.Abs"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
     %t1 = "tf.Acos"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
     %t2 = "tf.Acosh"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>

func.func @broadcast_mul0(%arg0: tensor<5x7xf32>, %arg1: tensor<7xf32>) -> tensor<5x7xf32> {
  %cst = arith.constant dense<[5, 7]> : tensor<2xi32>
  %0 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<7xf32>, tensor<2xi32>) -> tensor<5x7xf32>
  %1 = "tf.Mul"(%arg0, %0) : (tensor<5x7xf32>, tensor<5x7xf32>) -> tensor<5x7xf32>
  func.return %1 : tensor<5x7xf32>
  // CHECK: %[[V0:.*]] = "tf.Mul"(%arg0, %arg1) : (tensor<5x7xf32>, tensor<7xf32>) -> tensor<5x7xf32>

  (TF_LessEqualOp:$dest $arg0, $arg1), [], [(CopyAttrs $src, $dest)]>;

def LogicalNotOfGreaterEqual : Pat<
  (TF_LogicalNotOp:$src (TF_GreaterEqualOp $arg0, $arg1)),
  (TF_LessOp:$dest $arg0, $arg1), [], [(CopyAttrs $src, $dest)]>;

def LogicalNotOfLess : Pat<(
  TF_LogicalNotOp:$src (TF_LessOp $arg0, $arg1)),
  (TF_GreaterEqualOp:$dest $arg0, $arg1), [], [(CopyAttrs $src, $dest)]>;

// CHECK: func private @mul1(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tf.entry_function = {inputs = "mul1_y:0,mul1_x:0", outputs = "PartitionedCall:0"}}
// CHECK:   %[[MUL_0:.*]] = "tf.Mul"(%arg1, %arg0) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK:   return %[[MUL_0]] : tensor<1xf32>
// CHECK: }

func.func @testPwStreamResults(%arg0: tensor<f32>, %arg1: tensor<f32>) {
  "tf.PwStreamResults"(%arg0, %arg1) {names = ["foo", "bar"]} : (tensor<f32>, tensor<f32>) -> ()
  return
}

// -----
// CHECK-LABEL: func @test_ifrt_call
func.func @test_ifrt_call(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {

}
func testOr(t *testing.T, arg1, arg2, wantRes bool) {
	testShortCircuit(t, "OR", arg1, arg2, or_ssa, !arg1, wantRes)
}

func testShortCircuit(t *testing.T, opName string, arg1, arg2 bool, fn func(bool, bool) bool, wantRightCall, wantRes bool) {
	rightCalled = false
	got := fn(arg1, arg2)
	if rightCalled != wantRightCall {
		t.Errorf("failed for %t %s %t; rightCalled=%t want=%t", arg1, opName, arg2, rightCalled, wantRightCall)
	}