// CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64
  // CHECK-SAME: max_enqueued_batches = 10 : i64
  // CHECK-SAME: num_batch_threads = 2 : i64
  // CHECK-SAME: shared_name = "batch/"

        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
      } : (tensor<?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint8>
      %1 = "tf.UniformDequantize"(%0, %scales, %zps) {
        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
      } : (tensor<?x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?xf32>

  // CHECK: RemoteCall
  // CHECK-SAME: (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<i64>) -> tensor<!tf_type.string>
  %0 = "tf.RemoteCall"(%arg1, %arg0, %arg2) <{f = @__inference__next_func_3760}> {device = ""} : (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<i64>) -> tensor<!tf_type.string>
  return
}

// CHECK-NEXT:   [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
// CHECK-SAME:       used_by_host = false 
// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"([[KEY]], %arg0) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
// CHECK-SAME:    : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
//
module {

    %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
    return %0 : tensor<1x3x4x2xf32>
  }
}

// CHECK-LABEL: quantize_conv_fn
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x3x4x3xf32>

    %0:3 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<3xi64>) -> (tensor<i64>, tensor<i64>, tensor<i64>)
    return
  }
})";

// CHECK-NEXT:    "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    "tf.IfrtCall"(%arg0, [[ARRAYKEY]]) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
// CHECK-NEXT:    return
//

      tf_device.return

    return false;
  }
  if (AttrType content = mlir::dyn_cast<AttrType>(attr)) {
    *out_attr = content;
    return true;
  } else {
    return false;
  }
}

// Changes a DenseIntElementsAttr **containing I64** elements to an I32 Vector.
bool DenseI64AttrToI32Vector(const DenseIntElementsAttr& dense_attr,
                             std::vector<int32_t>* out_vec);

// Gets boolean from composite attrs if it exists.

func.func @gather(%indices: tensor<?xi32>,
             %resource: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
  %0 = "tf.ResourceGather"(%resource, %indices) {batch_dims = 0 : i64, device = "/device:CPU:0", validate_indices = true}: (tensor<*x!tf_type.resource>, tensor<?xi32>) -> (tensor<*xi32>)
  func.return %0 : tensor<*xi32>
}