"tf_device.cluster"() ({
      %0 = "tf.RecvTPUEmbeddingActivations"() {config = "test_config_recv_embedding"} : () -> tensor<512x256xf32>
      "tf.SendTPUEmbeddingGradients"(%0) {N = 1 : i64, NN = 0 : i64, config = "test_config_send_embedding", operandSegmentSizes = array<i32: 1, 0>} : (tensor<512x256xf32>) -> ()
      "tf.A"() {_xla_outside_compilation = "cluster1"} : () -> ()
      tf_device.return

  %3:4 = "tf.While"(%0, %1, %0, %2) {_lower_using_switch_merge = true, _num_original_outputs = 4 : i64, _read_only_resource_inputs = [], body = @otherVariantWhileBody, cond = @otherVariantWhileCond, device = "", is_stateless = true, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>)

  // Check bytes come back in the right order
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 16909060]> : tensor<4xi32> } : () -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

func.func @i64() -> tensor<4xi64> {
  // CHECK-LABEL: @i64
  // CHECK: value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64> } : () -> tensor<4xi64>

  // Check bytes come back in the right order
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 16909060]> : tensor<4xi32> } : () -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

func.func @i64() -> tensor<4xi64> {
  // CHECK-LABEL: @i64
  // CHECK: value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64> } : () -> tensor<4xi64>

  %cst = "tf.Const"() {device = "", value = dense<[1]> : tensor<1xi32>} : () -> tensor<1xi32>
  // idx 2 should be a compile time constant
  %0 = "tf.GatherV2"(%arg0, %arg1, %cst) {batch_dims = 1: i64} : (tensor<16x4xf32>, tensor<16xi32>, tensor<1xi32>) -> tensor<16xf32>

  // Just to introduce an extra use for %cst.
  %1 = "tf.AddV2"(%cst, %cst) {device = ""} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>

  // CHECK:   tf._TPUDeviceOrdinalPlaceholder
  %0:2 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*xf32>) {n = 2: i32} {
    %1 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 0} : () -> tensor<i64>
    tf_device.return %1 : tensor<i64>
  }
  func.return
}


// CHECK-LABEL: func @replicate_resource_var_arg_shape
// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: tensor<*x!tf_type.resource>, %[[ARG_1:[a-z0-9]*]]: tensor<*x!tf_type.resource>)

      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
      %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
      return %2 : tensor<?x2xf32>
    }

    %mul = "tf.Mul"(%input, %weight_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %mul : tensor<*xf32>
  }

  // Note that input i64 type is also supported by this.
  parameters[
    {"quantized_ops": ["Gather"]}
  ]
  func.func @GenerateQuantizedFunctionName(${quantized_ops})(

        Sout = [#tf_type.shape<?x10>], dim_args_spec = [],
        disabled_checks = [], function_list = [],
        has_token_input_output = false, module = "", platforms = [],
        version = 5 : i64
    }> {
        _entry_function = @composite_dot_general_fn_1,
        _original_entry_function = "composite_dot_general_fn_1",
        _tfl_quant_trait = "fully_quantizable",

  // CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) <{device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"}> {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> ()