return
      }
    }
  )mlir");

  FailureOr<SmallVector<AssetFileDef>> asset_file_defs =
      ConvertAssetArgs(*module_op);

  EXPECT_TRUE(succeeded(asset_file_defs));
  EXPECT_THAT(*asset_file_defs, SizeIs(1));

    return failure();
  }
  return static_cast<int32_t>(value);
}

FailureOr<SmallVector<int32_t>> CastI64ArrayToI32(
    const ArrayRef<int64_t> int64_array) {
  SmallVector<int32_t> int32_array{};
  int32_array.reserve(int64_array.size());

  for (const int64_t i64 : int64_array) {
    FailureOr<int32_t> cast_i32 = CastI64ToI32(i64);
    if (failed(cast_i32)) return failure();

namespace mlir {

// Set layouts attribute of tf.InfeedDequeueTuple ops.
bool SetTPUInfeedLayout(OwningOpRef<ModuleOp>& mlir_module);

// Try to determine the right TPU infeed layout.
FailureOr<Attribute> GetTPUInfeedLayout(ArrayRef<Type> types,
                                        OpBuilder& rewriter);

}  // namespace mlir

      : OpRewritePattern<RootOpT>(context, /*benefit=*/300) {}

 private:
  // Collects all candidate ops for quantization, which are the
  // `dequantize_op`'s users.
  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
      DequantizeOpT dequantize_op) const {
    auto users = dequantize_op->getResult(0).getUsers();
    return SmallVector<Operation*>(users.begin(), users.end());
  }

constexpr llvm::StringRef kCalledIndexAttrName = "called_index";
constexpr llvm::StringRef kCalledFuncAttrName = "called_func";

// Converts `called_func` attributes in custom call ops back to `called_index`.
FailureOr<ArrayAttr> DesymbolizeCustomCallCalledIndex(ModuleOp module) {
  Builder builder(module.getContext());

  SmallVector<Attribute> function_list;
  llvm::DenseMap<SymbolRefAttr, int> called_indexes;

// Returns a failure() on error, otherwise returns true to indicate that it
// reached convergence, false otherwise.
// If input shapes are provided, first refines the `main` function using
// InferShapeForFunction.
FailureOr<bool> InferModuleShape(ModuleOp module, int64_t max_iterations = 10,
                                 ArrayRef<TypeID> ops_to_skip = {},
                                 ArrayRef<ArrayRef<int64_t>> input_shapes = {});

// dictionary attribute.
//
// If the op does not represent a TF function call, returns nullptr.
// Otherwise, if the op does not have `caller_name`, returns failure.
FailureOr<SymbolRefAttr> GetTfFuncCustomCallFuncName(
    stablehlo::CustomCallOp op);

}  // namespace TF
}  // namespace mlir

#include "mlir/Support/LogicalResult.h"  // from @llvm-project

namespace mlir::quant::tensorflow {

// GetDenseAttrFromTensorProtoAttr returns DenseElementsAttr from tensor proto.
FailureOr<mlir::DenseElementsAttr> GetDenseAttrFromTensorProtoAttr(
    llvm::StringRef mangled_tensor_proto, TensorType result_tensor_type);

// Check if a type is TF qint type.
bool IsTFQintType(Type type);

                                 << T::getOperationName() << ").\n");
    return failure();
  }
}

FailureOr<int32_t> CastI64ToI32(int64_t value);

// Tries to cast an array of int64 to int32. If any of the element in the
// array is not in the range of int32, returns failure().
FailureOr<SmallVector<int32_t>> CastI64ArrayToI32(
    ArrayRef<int64_t> int64_array);

//   "tf_saved_model.index_path", if the bound input is attached to the
//   `tf_saved_model::AssetOp`.
// * Strips the "assets/" prefix of the filename when setting it to
//   `AssetFileDef`.
FailureOr<SmallVector<tensorflow::AssetFileDef>> ConvertAssetArgs(
    ModuleOp module_op);

}  // namespace mlir::quant