func.return %0 : tensor<*xf32>
}

// CHECK: func @xla_gather
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi64>}> : () -> tensor<2xi64>
// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<1x1xi64>}> : () -> tensor<1x1xi64>

  func.return %2 : tensor<4xf32>

// CHECK-NEXT: %[[CST:.*]] = arith.constant dense<1.000000e+00> : tensor<4xf32>
// CHECK-NEXT: %[[MUL:.*]] = tfl.mul %arg0, %[[CST]] {fused_activation_function = "NONE"} : tensor<4xf32>
// CHECK-NEXT: %[[CUSTOM_1:.*]] = "tfl.custom_tf"(%[[MUL]], %[[CST]]) ({
// CHECK-NEXT: ^bb0(%arg1: tensor<4xf32>, %arg2: tensor<4xf32>):

    // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense<"assets/table.txt"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
    // CHECK: [[HASHTABLE:%.+]] = "tf.HashTableV2"()
    // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE]], [[CST]])
    func.return
  }
}

// -----

module attributes {tf_saved_model.semantics} {

    func.return %arg0 : tensor<*xi1>
  }
  func.func @"tpu_subgraph"(%arg0: tensor<*xi64>) -> tensor<*xi1> {
    %cst = "tf.Const"() {value = dense<0> : tensor<i64>} : () -> tensor<i64>
    %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>

  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %[[FQ:.*]] = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2)
// CHECK:  %[[T1:.*]] = "tf.Transpose"(%[[FQ]], %cst)
// CHECK:  %[[T2:.*]] = "tf.Transpose"(%[[FQ]], %cst)
}

// CHECK-LABEL: fakeQuantFollowedByReshape

// CHECK: "tfl.fully_connected"(%arg0, %arg1, [[NONE]])
// CHECK-SAME: (tensor<40x37xf32>, tensor<40x37xf32>, none) -> (tensor<40x40xf32>, tensor<40x40xf32>)
  %cst = "tfl.no_value"() {value = unit} : () -> none
  %0:2 = "tfl.fully_connected"(%arg0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, none) -> (tensor<40x40xf32>, tensor<40x40xf32>)

  func.return %7 : tensor<1x3x2x2xf32>
}

// CHECK: func @tpu_conv(%[[ARG0:.*]]: tensor<1x3x4x3xf32>)
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense_resource<__elided__> : tensor<2x3x3x2xbf16>}> {device = ""} : () -> tensor<2x3x3x2xbf16>
// CHECK: %[[cast:.*]] = "tf.Cast"(%[[cst]]) <{Truncate = false}> : (tensor<2x3x3x2xbf16>) -> tensor<2x3x3x2xf32>
// CHECK: %[[conv:.*]] = "tf.Conv2D"(%[[ARG0]], %[[cast]])

		t.Run(tt.name, func(t *testing.T) {
			cst := NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("X-Proto", r.Proto)
			}))

			switch tt.name {
			case "http2":
				cst.EnableHTTP2 = true
				cst.StartTLS()
			default:
				cst.Start()
			}

			defer cst.Close()

			res, err := cst.Client().Get(cst.URL)
			if err != nil {

      %cst = "tf.Const"() {value = dense<1024> : tensor<i32>} : () -> tensor<i32>
      %0:2 = "tf_device.launch"() ({
            %1 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>
            %2 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>

}

// -----

module {
  func.func @quantize_sharded_weights_with_xladot(%arg0: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xbf16> {
    %cst = "tf.Const"() {device = "", value = dense<1.000000e+01> : tensor<512x512xf32>} : () -> tensor<512x512xf32>