// CHECK-SAME: , asymmetric_quantize_inputs = true
// CHECK: return %[[mm:.*]]

// PerTensor: %[[w:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<512x12x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>,
// PerTensor: %[[mm:.*]] = "tfl.batch_matmul"(%arg0, %[[w]]) <{adj_x = false, adj_y = false
// PerTensor-SAME: , asymmetric_quantize_inputs = true
// PerTensor: return %[[mm:.*]]
}

// disabled by having `_quantization_method = "no_quantization { }"`.
// DISABLE-ALL-DOT-GENERAL-SAME: _entry_function = @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: _original_entry_function
// DISABLE-ALL-DOT-GENERAL-SAME: _quantization_method = "no_quantization { }"
// DISABLE-ALL-DOT-GENERAL-SAME: _tfl_quant_trait = "fully_quantizable"

// DISABLE-ALL-DOT-GENERAL: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>

// CHECK-LABEL: func @dot_int4
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot_int4(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
  // CHECK: %[[cst:.*]] = stablehlo.constant
  // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
  // CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>

  // FB-SAME:  -8.0, -9.0, -10.0, -11.0, -12.0, -13.0, -14.0, -15.0, -16.0,
  // FB-SAME:  -17.0, -18.0, -19.0, -20.0, -21.0, -22.0, -23.0, -24.0, -25.0,
  // FB-SAME:  -26.0, -27.0, -28.0, -29.0, -30.0, -31.0, -32.0, -33.0, -34.0,
  // FB-SAME:  -35.0, -36.0, -37.0, -38.0, -39.0 ],
  // FB-NEXT:        max: [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0,

  func.return %3 : tensor<4x8xf32>
}

// CHECK-LABEL:   func @select(
// CHECK-SAME:                 %[[VAL_0:.*]]: tensor<2xi1>,
// CHECK-SAME:                 %[[VAL_1:.*]]: tensor<2xi32>,
// CHECK-SAME:                 %[[VAL_2:.*]]: tensor<2xi32>) -> tensor<2xi32> {

// TODO(b/316226111): the printer may not guarantee the same order of fields. Rewrite the checks to be less sensitive to proto serialization formats.
// -----
// Non-SPMD: one input and one output
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
// CHECK-NEXT:  "tf.IfrtCall"(%arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
// CHECK-SAME:       (tensor<1x3xf32>) -> tensor<1x3xf32>

  // CHECK-SAME: %arg21: tensor<!tf_type.resource<tensor<640xf32>>>
  // CHECK-SAME: %arg22: tensor<!tf_type.resource<tensor<640xf32>>>
  // CHECK-SAME: %arg23: tensor<!tf_type.resource<tensor<512x640xf32>>>
  // CHECK-SAME: %arg24: tensor<!tf_type.resource<tensor<640x640xf32>>>
  // CHECK-SAME: %arg25: tensor<!tf_type.resource<tensor<4096xf32>>>
  // CHECK-SAME: %arg26: tensor<!tf_type.resource<tensor<640x4096xf32>>>

  // CHECK-SAME: _XlaSharding = ""
  // CHECK-SAME: partition_dims = []
  // CHECK-DAG:  "tf.AssignVariableOp"([[ARG0]], [[OUTPUT]]#0)
  // CHECK-DAG:  "tf.AssignVariableOp"([[ARG1]], [[OUTPUT]]#1)
  "tf.AssignVariableOp"(%0, %2) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
  func.return
}

// CHECK-LABEL: func @read_write_packed_resource

// RUN: tf-opt -split-input-file -verify-diagnostics -xla-legalize-tf-collective -xla-legalize-tf %s | FileCheck %s


// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 0
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_size = 2
// CHECK-LABEL: func @all_reduce_cross_replica
func.func @all_reduce_cross_replica(%input: tensor<f32>) -> tensor<f32> {

  }

// CHECK-LABEL: func @conv_with_single_layer

// CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%arg0
// CHECK-SAME: f = @quantize_i8
// CHECK: %[[conv_quant:.*]] = "tf.PartitionedCall"(%[[quantize]]
// CHECK-SAME: f = @quantized_conv2d_with_bias_and_relu6_float_output_fn_0