LLVM_DEBUG(llvm::dbgs() << "Zero point value should be a tensor type. Got: "
                            << zero_points_type << ".\n");
    return failure();
  }

  if (Type zero_points_element_type = zero_points_type.getElementType();
      !mlir::isa<IntegerType>(zero_points_element_type)) {
    LLVM_DEBUG(llvm::dbgs() << "Zero point should be an integer type. Got: "
                            << zero_points_element_type << ".\n");

        the layer's dtype. If you had a model that used `float64`, it will
        probably silently use `float32` in TensorFlow2, and a warning will be
        issued that starts with Layer "layer-name" is casting an input tensor
        from dtype float64 to the layer's dtype of float32. To fix, either set
        the default dtype to float64 with

  TF_OperationDescription* init(string type) {
    // Construct op_name to match the name used by REGISTER_OP in the
    // ATTR_TEST_REGISTER calls above.
    string op_name = "CApiAttributesTestOp";
    if (type.find("list(") == 0) {
      op_name += "List";
      type = type.replace(0, 5, "");
      type = type.replace(type.size() - 1, 1, "");
    }
    op_name += type;
    return TF_NewOperation(

  // An additional sharding is added at the end to account for token result.
  // Proto debug string:
  //   type: TUPLE
  //   tuple_shardings {
  //     type: MAXIMAL
  //     tile_assignment_dimensions: 1
  //     tile_assignment_devices: 0
  //   }
  //   tuple_shardings {
  //     type: MAXIMAL
  //     tile_assignment_dimensions: 1
  //     tile_assignment_devices: 0
  //   }

}

// Return true if the passed quantized type is unsigned.
bool QuantizedTypeIsUnsigned(Type type) {
  return TypeSwitch<Type, bool>(type)
      .Case<mlir::TF::Qint8Type>([](Type) { return false; })
      .Case<mlir::TF::Qint16Type>([](Type) { return false; })
      .Case<mlir::TF::Qint32Type>([](Type) { return false; })
      .Case<mlir::TF::Quint8Type>([](Type) { return true; })

        llvm::dyn_cast<TF::ConstOp>(fold_results[0].getDefiningOp()).getValue(),
        &tensor));
    xla::Shape xla_shape;
    TF_RETURN_IF_ERROR(tensorflow::TensorShapeToXLAShape(
        tensor.dtype(), tensor.shape(), &xla_shape));
    xla::PjRtClient::HostBufferSemantics host_buffer_semantics =
        xla::PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes;
    TF_ASSIGN_OR_RETURN(
        auto buffer,

  // CHECK: %[[RESULT:.*]] = "tf.Maximum"(%arg0, %[[CST]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0", dtype = i32} : (tensor<4xi32>, tensor<i32>) -> tensor<4xi32>
  // CHECK: return %[[RESULT]]
  %cst_0 = arith.constant dense<0> : tensor<i32>
  %0 = "tf.Maximum"(%arg0, %cst_0) {device = "/job:localhost/replica:0/task:0/device:GPU:0", dtype = i32} : (tensor<4xi32>, tensor<i32>) -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

//   args {
//     dtype: DT_INT32
//     shape {
//       dim {
//         size: 8
//       }
// .   }
//     kind: PARAMETER
//     sharding {
//       type: MAXIMAL
//       tile_assignment_dimensions: 1
//       tile_assignment_devices: 0
//     }
//   }
//   retvals {
//     sharding {
//       type: MAXIMAL
//       tile_assignment_dimensions: 1

  // sharding.
  // The split sharding is:
  //   type: OTHER
  //   tile_assignment_dimensions: 2
  //   tile_assignment_dimensions: 1
  //   tile_assignment_devices: 0
  //   tile_assignment_devices: 1
  // Serialized string:
  //   "\08\03\1A\02\02\01\22\02\00\01"
  // The manual sharding is:
  //   type: MANUAL
  // Serialized string:
  //   "\08\04"

    save.save(
        model,
        self._input_saved_model_path,
        signatures=model.matmul.get_concrete_function(
            tensor_spec.TensorSpec(
                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
            )
        ),
    )

    def data_gen() -> repr_dataset.RepresentativeDataset:
      for _ in range(50):