axis_op);
}

ArrayRef<int64_t> GetRankedTensorShape(Value value) {
  return mlir::cast<RankedTensorType>(value.getType()).getShape();
}

Value SliceRankedTensor(OpBuilder* builder, Value input,
                        ArrayRef<int64_t> begin_shape,
                        ArrayRef<int64_t> begin_values,

  DenseIntElementsAttr iota_const_attr;
  if (!matchPattern(iota, m_Constant(&iota_const_attr))) return false;

  auto iota_type = iota_const_attr.getType();
  auto iota_shape = iota_type.getShape();
  auto reduce_dim = (*dimensions.value_begin<APInt>()).getSExtValue();
  if (reduce_dim < 0) reduce_dim += iota_type.getRank();

  auto index =

  auto ranked_ty = mlir::dyn_cast_or_null<mlir::RankedTensorType>(ty);
  if (!ranked_ty) return {};

  int64_t rank = ranked_ty.getRank();
  llvm::SmallVector<int64_t, 4> dims = llvm::to_vector<4>(ranked_ty.getShape());
  auto encoding = mlir::dyn_cast_or_null<mlir::mhlo::TypeExtensionsAttr>(
      ranked_ty.getEncoding());
  if (encoding && !encoding.getBounds().empty()) {
    for (int64_t dim = 0; dim < rank; ++dim) {

      return op_->emitRemark()
             << "skipping legalization due to unsupported type "
             << operand.getType();
    }

    auto shape_or = expr.GetShape();
    if (!shape_or.ok()) {
      return op_->emitRemark()
             << "failed to get shape for expression. " << expr.HumanString();
    }

    tensors.emplace_back(

// Revert the permutation applied in `type`.
static mlir::ShapedType ReversePermuteShapedType(
    mlir::ShapedType type, ArrayRef<int64_t> permutation) {
  if (!type.hasRank()) return type;

  auto shape = type.getShape();
  SmallVector<int64_t, 4> new_shape(shape.size());

  for (int i = 0; i < permutation.size(); ++i) {
    int64_t index = permutation[i];
    assert(index < shape.size());
    new_shape[index] = shape[i];

        aqtype.getZeroPoints(), new_out_quant_dim, /*narrow_range=*/true);
    auto new_tensor_type = RankedTensorType::getChecked(
        transpose_op.getLoc(), transpose_op.getType().getShape(), new_qtype);
    rewriter.setInsertionPointAfter(transpose_op);
    auto new_q_op = rewriter.create<quantfork::QuantizeCastOp>(
        transpose_op.getLoc(), new_tensor_type, q_op.getArg());

      ".getShape()[" # dim # " ] == " # size>]>;

// Returns true if the n-th operand is ranked and has a dimension length <=
// size at the rank dim.
class TFL_OperandDimIsAtMost<int n, int dim, int size> : And<[
  TFL_OperandIsRankedAndHasDimPred<n, dim>,
  CPred<"$_op.getOperand(" # n # ").getType().cast<ShapedType>()"
      ".getShape()[" # dim # " ] <= " # size>]>;

  if (!elem_type.has_value()) {
    return stack.emitOpError("cannot infer element shape of stack");
  }
  OpBuilder builder(stack);
  Value buffer;
  if (failed(cutil::CreateInitBufferValue(
          elem_type->getShape(), stack.getMaxSize(), stack,
          elem_type->getElementType(), builder, &buffer))) {
    return failure();
  }
  auto size_var_type = GetSizeVarType(builder);
  auto var_type = RankedTensorType::get(

  }
  elt_type = TensorFlowRefType::get(elt_type);
  if (RankedTensorType tensor_type = mlir::dyn_cast<RankedTensorType>(type)) {
    *output_type = RankedTensorType::get(tensor_type.getShape(), elt_type);
  }
  *output_type = UnrankedTensorType::get(elt_type);
  return absl::OkStatus();
}

Status MlirAbstractOp::Create(ArrayRef<Value> operands,

  if (!ranked_type)
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output types must be "
              "ranked tensors.";
  ArrayRef<int64_t> in_shape = ranked_type.getShape();
  if (in_shape.empty() || in_shape[0] < 0) {
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output shapes must "