// CHECK-PER-TENSOR: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]])

// get a data token.

// CHECK-LABEL: func @unconnected_while_body
func.func @unconnected_while_body(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>, %arg3: tensor<f32>) -> (!tf_res, !tf_res, tensor<f32>, tensor<f32>) {
  // CHECK: tf_executor.graph
  %graph:4 = tf_executor.graph {
    // CHECK: %[[C0:.*]] = {{.*}}AssignVariableOp
    // CHECK: %[[C1:.*]] = {{.*}}AssignVariableOp
    // CHECK: VarHandleOp

  // CHECK: tf.Const
  // CHECK: tf.Mul
  // CHECK: tf.Mul
  // CHECK: tf.Div
  // CHECK: tf.Mul
  %0 = "tf.RsqrtGrad"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

// CHECK-LABEL: func @sqrt_grad_unranked
// CHECK-SAME: (%[[ARG0:.*]]: tensor<*xcomplex<f32>>, %[[ARG1:.*]]: tensor<*xcomplex<f32>>)

  func.return
}

// CHECK: [[CONST:%.*]] = "tf.Const"
// CHECK: [[VAR:%.*]] = "tf.VarHandleOp"
// CHECK: "tf_device.cluster"() ({
// CHECK:   "tf.opA"() ({
// CHECK:     "tf.AssignAddVariableOp"([[VAR]], [[CONST]])
// CHECK:    })
// CHECK:    "tf.opB"()
// CHECK:    tf_device.return
// CHECK:  })
// CHECK-SAME: _replication_info = "foo"
// CHECK: "tf.AssignAddVariableOp"([[VAR]], [[CONST]])

// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }

// CHECK-LABEL: private @composite_dot_general_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: return %[[DOT_GENERAL:.*]] : tensor<1x1x64xf32>
// CHECK: }

// -----

// CHECK-LABEL: @dot_general_with_bias_same_shape_fn(

// parameterized type or function value.
func (check *Checker) rawExpr(T *target, x *operand, e ast.Expr, hint Type, allowGeneric bool) exprKind {
	if check.conf._Trace {
		check.trace(e.Pos(), "-- expr %s", e)
		check.indent++
		defer func() {
			check.indent--
			check.trace(e.Pos(), "=> %s", x)
		}()
	}

	kind := check.exprInternal(T, x, e, hint)

	if !allowGeneric {
		check.nonGeneric(T, x)
	}

    // CHECK: [[GRAD_SQUARE:%.*]] = "tf.Square"([[GRAD]])
    // CHECK: [[V_DELTA:%.*]] = "tf.Mul"([[ONE_MINUS_BETA2]], [[GRAD_SQUARE]])
    // CHECK: [[NEW_V:%.*]] = "tf.AddV2"([[BETA2_OLD_V]], [[V_DELTA]])
    // CHECK: [[ALPHA_NEW_M:%.*]] = "tf.Mul"([[ALPHA]], [[NEW_M]])
    // CHECK: [[SQRT_NEW_V:%.*]] = "tf.Sqrt"([[NEW_V]])
    // CHECK: [[SQRT_NEW_V_EPSILON:%.*]] = "tf.AddV2"([[SQRT_NEW_V]], [[EPSILON]])

  // CHECK:          "tf.C"
  // CHECK:          "tf.D"
  // CHECK:        "tf.AB"
  // CHECK:        "tf.IfRegion"
  // CHECK:          "tf.E"
  // CHECK:          "tf.F"
  // CHECK:        "tf.IfRegion"
  // CHECK:          "tf.Const"() <{value = dense<1.000000e+00> : tensor<f32>}> : () -> tensor<f32>
  // CHECK:          "tf.Const"() <{value = dense<3.000000e+00> : tensor<f32>}> : () -> tensor<f32>

  func.return %4: tensor<2x4xf32>

// CHECK-NEXT: %[[q:.*]] = "tfl.quantize"(%arg0)
// CHECK-NEXT: %[[dq:.*]] = "tfl.dequantize"(%[[q]])
// CHECK-NEXT: %[[q_0:.*]] = "tfl.quantize"(%arg1)
// CHECK-NEXT: %[[dq_0:.*]] = "tfl.dequantize"(%[[q_0]])
// CHECK-NEXT: %[[c:.*]] = "tfl.concatenation"(%[[dq]], %[[dq_0]])
// CHECK-NEXT: %[[q_1:.*]] = "tfl.quantize"(%[[c]])
// CHECK-NEXT: %[[dq_1:.*]] = "tfl.dequantize"(%[[q_1]])

// CHECK: "tf.AssignVariableOp"(%[[CARG1]], %[[UPDATE1]])
// CHECK: %[[READ2:.*]] = "tf.ReadVariableOp"(%[[CARG2]]) : (tensor<!tf_type.resource<tensor<5x3xf32>>>) -> tensor<5x3xf32>
// CHECK: %[[UPDATE2:.*]] = "tf.XlaDynamicUpdateSlice"(%[[READ2]],
// CHECK: "tf.AssignVariableOp"(%[[CARG2]], %[[UPDATE2]])
// CHECK: return %[[CARG0]]

// -----

// Tests (Stateful)PartitionedCall op with private callee function.

// CHECK-LABEL: func @main