}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_2(%arg0: tensor<1x1x1x8xf32>) -> tensor<1x1x1x4xf32>

// CHECK-LABEL: avg_pool2d_3
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x8xf32>, tensor<4xi32>) -> tensor<1x1x8x1xf32>
// CHECK{LITERAL}: %cst_0 = arith.constant dense<[[0, 0], [0, 0], [1, 1], [0, 0]]> : tensor<4x2xi32>

  %cst = "tf.Const"() {value = dense<5.000000e+00> : tensor<f32>} : () -> tensor<f32>
  %0 = "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster_tpu_function", device = ""} : () -> tensor<!tf_type.string>

			GotConn: func(info httptrace.GotConnInfo) {
				rt.conn = info.Conn.(*transportDialTesterConn)
			},
		})
		req, _ := http.NewRequestWithContext(ctx, "POST", dt.cst.ts.URL, pr)
		req.Header.Set("Content-Type", "text/plain")
		rt.res, rt.err = dt.cst.tr.RoundTrip(req)
		dt.t.Logf("RoundTrip %v: done (err:%v)", rt.roundTripID, rt.err)
		close(rt.done)
	}()
	return rt
}

  %cst = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
  // Predicate should be inlined.
  // CHECK-NEXT: tf_mlrt.predicate
  // CHECK-NEXT: mlrt.cond
  // CHECK-NEXT: tf_mlrt.predicate

  // CHECK-NEXT: mlrt.while

func.func @invalid_non_mhlo() -> ( tensor<128x32x4xcomplex<f32>> {mhlo.sharding = ""}) {
  // expected-error @+1 {{Could not legalize op: tf.Const}}
  %cst = "tf.Const"() {value = dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>} : () -> tensor<128x32x4xcomplex<f32>>
  return %cst : tensor<128x32x4xcomplex<f32>>
}

// -----

func.func @invalid_mixed_mhlo() -> (tensor<8x64x128xcomplex<f32>> {mhlo.sharding = ""}) {

// CHECK-NEXT: "func.return"([[res]]) : (tensor<*xf32>) -> ()


// -----

// CHECK-LABEL: same_scale_test
func.func @same_scale_test(%arg0: tensor<*xf32>) -> tensor<*xf32> {
  %cst = arith.constant dense<[-1, 144]> : tensor<2xi32>
  %0 = "quantfork.qcast"(%arg0) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 5.000000e-02:-10>>

    %0 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsite("test@conv"("Conv2D") at "QuantizationUnit(\12\06Conv2D\1a\04conv)"))

    return %3 : tensor<6x5xf32>
  }
  // CHECK: %[[CST:.*]] = "tf.Const"()
  // CHECK-NEXT: %[[CALL1:.*]] = "tf.XlaCallModule"(%[[ARG0]], %[[CST]])
  // CHECK-SAME: _entry_function = @composite_dot_general_fn_1
  // CHECK-SAME: _tfl_quant_trait = "fully_quantizable"
  // CHECK-NOT: "tf.XlaCallModule"

    %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
    %0 = "quantfork.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>

  // expected-remark@+1 {{failed to create tf2xla kernel: INVALID_ARGUMENT: NodeDef missing attrs 'seed2', 'seed' from}}
  %cst = "tf.Const"() {value = dense<[32, 12, 12, 64]> : tensor<4xi32>} : () -> tensor<4xi32>
  %0 = "tf.RandomUniform"(%cst) {} : (tensor<4xi32>) -> tensor<32x12x12x64xf32>
  // CHECK: return %1 : tensor<32x12x12x64xf32>
  func.return %0 : tensor<32x12x12x64xf32>
}
}
// -----