limitations under the License.
==============================================================================*/

#include <algorithm>
#include <memory>
#include <vector>

#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/UseDefLists.h"  // from @llvm-project

    // This occurs when doing qi16 convolution, as bias is passed as a
    // non-quantized int64
    if (allTypesFp || allTypesQuantizedOrInt) return failure();

    Location loc = op->getLoc();
    SmallVector<Value> newOperands;
    newOperands.reserve(op->getNumOperands());
    for (auto operand : op->getOperands()) {
      if (QuantizedType::getQuantizedElementType(operand.getType())) {

  ArrayRef<int64_t> filter_shape = filter_value_attr.getShapedType().getShape();
  SmallVector<int8_t> filter_constant_values{
      filter_value_attr.getValues<int8_t>()};
  SmallVector<int8_t> new_filter_constant_values(filter_constant_values.size(),
                                                 0);
  SmallVector<int64_t, 4> transpose_dims;
  if (is_depthwise) {
    transpose_dims = {2, 0, 1, 3};
  } else {

  // Parse optional argument list (control dependencies only).
  SmallVector<OpAsmParser::UnresolvedOperand, 4> op_infos;
  if (parser.parseOperandList(op_infos, OpAsmParser::Delimiter::OptionalParen))
    return failure();
  if (!op_infos.empty()) {
    SmallVector<Type, 2> types(op_infos.size(), control_type);
    if (parser.resolveOperands(op_infos, types, loc, result.operands))

#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_XLA_REWRITE_UTIL_H_

#include <optional>
#include <string>
#include <utility>

#include "absl/strings/string_view.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"

#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONSTANT_FOLD_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONSTANT_FOLD_H_

#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/PatternMatch.h"  // from @llvm-project

namespace mlir {
namespace quant {

  // num_cores_per_replica x num_replicas. i-th row holds the operands for i-th
  // core. j-th column holds the operands for j-th replica.
  llvm::SmallVector<llvm::SmallVector<Value, 4>, 4>
      operands_per_replica_per_core(num_cores_per_replica);

  // Collect all operands in the 2D matrix.
  for (auto operand : replicated_input.getInputs()) {

    getUsedValuesDefinedAbove(island_op.getBody(), operands);

    SmallVector<Type, 16> func_operand_types;
    func_operand_types.reserve(operands.size());
    for (Value operand : operands)
      func_operand_types.push_back(operand.getType());

    // Function results are the yield operands
    SmallVector<Type, 16> func_result_types;
    for (Value operand : island_op.GetYield().getOperands())

  auto return_operands = llvm::to_vector<4>(return_op.getOperands());
  auto result_types = llvm::to_vector<4>(return_op.getOperandTypes());
  llvm::SmallVector<std::pair<int64_t, llvm::StringRef>, 4>
      output_only_resources;
  llvm::SmallVector<std::pair<int64_t, int64_t>, 4> input_output_alias;

  // Collect new return values for variable writes and either (a) output-only

  // each old output index in TPUExecute.
  llvm::SmallVector<int, 8> old_to_new_output_mapping;
  // The resources read by ReadVariableOps that are inputs to TPUExecute,
  // ordered by the input indices to TPUExecute.
  llvm::SmallVector<Value, 8> resources_read;
  // Operands for the new TPUExecuteAndUpdateVariables.
  llvm::SmallVector<Value, 8> new_operand_values;
};