include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// TODO(jpienaar): Move this somewhere general.
class GetI64ScalarElementsAttr<int value> :
  NativeCodeCall<"DenseElementsAttr::get<int64_t>(RankedTensorType::get({}, $_builder.getIntegerType(64)), " # value # ")">;

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op
     (TF_MapDatasetOp $input_dataset, $other_arguments, $f, $output_types,

      !IsQuantized(tensor)) {
    TF_ASSIGN_OR_RETURN(elem_type, GetCalibratedQuantizedType(tensor, builder));
  }

  if (tensor.shape.empty() && (is_constant || tensor.has_rank)) {
    return RankedTensorType::get({}, elem_type);
  }

  if (!tensor.shape_signature.empty()) {
    llvm::SmallVector<int64_t, 4> shape(tensor.shape_signature.begin(),
                                        tensor.shape_signature.end());

  return lhs.hasRank() && rhs.hasRank() && lhs.getRank() == rhs.getRank();
}

// Creates a compatible RankedTensorType where mismatched dimensions are
// replaced with dynamic sizes.
RankedTensorType GetCompatibleRankedTensorType(RankedTensorType lhs,
                                               RankedTensorType rhs) {
  assert(lhs.getRank() == rhs.getRank());
  llvm::SmallVector<int64_t, 4> dims;

      };

      SmallVector<Type, 4> output_types{
          value.getType(),
          RankedTensorType::get({}, rewriter.getF32Type()),
          RankedTensorType::get({}, rewriter.getF32Type()),
          RankedTensorType::get({effective_num_bins}, rewriter.getI64Type()),
      };

      // Insert custom aggregation op between operand and operator.

#include "tsl/platform/ml_dtypes.h"

namespace tensorflow {

using llvm::ArrayRef;
using llvm::SmallVector;
using mlir::Builder;
using mlir::DenseStringElementsAttr;
using mlir::ElementsAttr;
using mlir::RankedTensorType;
using mlir::ShapedType;
using mlir::Type;
using tensorflow::errors::InvalidArgument;

static TensorProto ConvertToProto(const Tensor& input_tensor,
                                  bool use_tensor_content = true) {

absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type,
    int value) {
  ShapedType dense_type = RankedTensorType::get(shaped_type.getShape(),
                                                shaped_type.getElementType());
  auto attr = CreateTypedAttr(dense_type, value);

  return rewriter->create<arith::ConstantOp>(loc, dense_type,

// of input type and no-op permutation.

std::pair<RankedTensorType, DenseIntElementsAttr> GetReformatTypeAndPermutation(
    int batch_dim, int feature_dim, int spatial_dim_start,
    int default_batch_dim, int default_feature_dim,
    int default_spatial_dim_start, int num_spatial_dims, RankedTensorType type,
    ConversionPatternRewriter& rewriter);

using mlir::DenseIntElementsAttr;
using mlir::failure;
using mlir::Location;
using mlir::LogicalResult;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::success;
using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ConstOp;

      auto global_tensor = LookupBoundInputOfType<GlobalTensorOp>(
          func, arg_number, symbol_table);
      if (!global_tensor) continue;

      auto arg_type = mlir::cast<RankedTensorType>(arg.getType());
      assert(arg_type.getRank() == 0);
      llvm::ArrayRef<TensorType> underlying_type =
          mlir::cast<TF::ResourceType>(arg_type.getElementType()).getSubtypes();

  for (auto input : to_graph.getOperands()) {
    Operation* input_op = input.getDefiningOp();
    if (input_op && input_op == from_graph.getOperation()) {
      auto input_type =
          mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
      if (input_type == nullptr || !input_type.hasStaticShape()) continue;
      // Quantized type does not support getSizeInBits.
      if (IsQUI8Type(input_type) || IsQI8Type(input_type)) {