"tf_device.cluster"() ({

    // CHECK: "tf.Const"() <{value = dense<(1.000000e+00,0.000000e+00)> : tensor<complex<f32>>}> : () -> tensor<complex<f32>>
    // CHECK-NOT: tf.ResourceApplyAdam

  // CHECK: return %[[BARG0]], %[[SUB]], %[[BARG2]]
  func.return %arg0, %sub : tensor<!tf_type.resource>, tensor<i32>
}
// CHECK: func @while_cond(%[[CARG0:.*]]: tensor<!tf_type.resource<tensor<5x3xf32>>>, %[[CARG1:.*]]: tensor<i32>, %[[CARG2:.*]]: tensor<!tf_type.resource<tensor<5x3xf32>>>)
func.func @while_cond(%arg0: tensor<!tf_type.resource>, %arg1: tensor<i32>) -> tensor<i32> {

  func.return
}
// CHECK: func @while_body(%[[BARG0:.*]]: tensor<f32>)
func.func @while_body(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) -> (tensor<*x!tf_type.resource<tensor<f32>>>) {
  // CHECK-NEXT: return %[[BARG0]]
  func.return %arg0 : tensor<*x!tf_type.resource<tensor<f32>>>
}
// CHECK: func @while_cond(%[[CARG0:.*]]: tensor<f32>)
func.func @while_cond(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32> {

		p := s.Prog(st)
		p.To.Type = obj.TYPE_MEM
		p.To.Reg = arg0
		p.From.Type = obj.TYPE_REG
		p.From.Reg = arg1

	case ssa.OpPPC64LoweredAtomicCas64,
		ssa.OpPPC64LoweredAtomicCas32:
		// MOVD        $0, Rout
		// LWSYNC
		// loop:
		// LDAR        (Rarg0), MutexHint, Rtmp
		// CMP         Rarg1, Rtmp
		// BNE         end
		// STDCCC      Rarg2, (Rarg0)
		// BNE         loop
		// MOVD        $1, Rout

// CHECK-LABEL: identity
// CHECK: %0 = "tf.Identity"(%arg0) : (tensor<10xi32>) -> tensor<*xi32>
// CHECK: return %arg0, %arg1, %arg2, %0
}

func.func @sharding(%arg0: tensor<10x10xi32>) -> (tensor<10x10xi32>) {
  %0 = "tf.MatMul"(%arg0, %arg0) {device = "", transpose_a = false, transpose_b = false} : (tensor<10x10xi32>, tensor<10x10xi32>) -> tensor<10x10xi32>

// CHECK: %[[dw:.*]] = "tfl.depthwise_conv_2d"(%arg0, %arg1, %arg2)
// CHECK: %[[q:.*]] = "tfl.quantize"(%[[dw]])
// CHECK: %[[dq:.*]] = "tfl.dequantize"(%[[q]])
// CHECK: return %[[dq]], %[[dq]]
}

// Series of values needing requantization -- first the args then the results
// of concatenation operations. concat(concat(arg2, arg0), concat(arg1, arg0)),
// concat(concat(arg2, arg0), arg3)). arg0 should be requantized twice --

		{name: "CMPU", argLength: 2, reg: gp2flags, asm: "CMPU", typ: "Flags"},   // arg0 compare to arg1
		{name: "CMPWU", argLength: 2, reg: gp2flags, asm: "CMPWU", typ: "Flags"}, // arg0 compare to arg1

		{name: "CMPconst", argLength: 1, reg: gp1flags, asm: "CMP", typ: "Flags", aux: "Int32"},     // arg0 compare to auxint

  %0 = "tf.Pack"(%arg0, %arg1) : (tensor<3x5xf32>, tensor<3x5xf32>) -> tensor<2x3x5xf32>
  func.return %0 : tensor<2x3x5xf32>
}

// CHECK-LABEL: func @square
func.func @square(%arg0: tensor<3xf32>) -> tensor<3xf32> {
  // CHECK: "tf.Mul"(%arg0, %arg0)
  %1 = "tf.Square"(%arg0) : (tensor<3xf32>) -> tensor<3xf32>
  func.return %1 : tensor<3xf32>
}

  %0 = "tf.TPUReplicatedInput"(%arg0, %arg0) : (!tf_res, !tf_res) -> !tf_res
  %1 = "tf.TPUReplicatedInput"(%arg0, %arg0) : (!tf_res, !tf_res) -> !tf_res
  %2 = "tf.TPUReplicatedInput"(%arg0, %arg0) : (!tf_res, !tf_res) -> !tf_res
  %3 = "tf.ReadVariableOp"(%0) {_xla_compile_device_type = "TPU", _replication_info = "a"} : (!tf_res) -> tensor<f32>

}

// -----

func.func @batchMatMulMatrix(%arg0: tensor<4x5xf32>, %arg1: tensor<5x6xf32>) -> tensor<4x6xf32> {
  %0 = "tf.BatchMatMul"(%arg0, %arg1) : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32>
  func.return %0 : tensor<4x6xf32>

  // CHECK-LABEL: batchMatMulMatrix