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 "

    llvm::ArrayRef<int64_t> shape_ref =
        mlir::cast<TensorType>(tensor_attr.getType()).getShape();
    if (mlir::failed(check_shape(shape_ref))) return std::nullopt;

    shape = std::vector<int32_t>(shape_ref.begin(), shape_ref.end());
  } else if (type.hasRank()) {
    llvm::ArrayRef<int64_t> shape_ref = type.getShape();
    if (mlir::failed(check_shape(shape_ref))) return std::nullopt;

    auto output_uniform_quantized_tensor_type = RankedTensorType::getChecked(
        rewriter.getUnknownLoc(),
        /*shape=*/
        mlir::cast<TensorType>(conv_output_value.getType()).getShape(),
        output_uniform_quantized_type);

    SmallVector<Type> new_conv_output_types = {
        output_uniform_quantized_tensor_type};
    auto new_conv_op_with_output_type =

                VALIDATE.getPhase(),
                INITIALIZE.getPhase(),
                PROCESS_RESOURCES.getPhase(),
                COMPILE.getPhase(),
                TEST.getPhase(),
                PROCESS_TEST_RESOURCES.getPhase(),
                PACKAGE.getPhase(),
                "BEER",
                INSTALL.getPhase());

    %2 = "tf.PartitionedCall"(%1) {config = "", config_proto = "", executor_type = "", f = @get_shape} : (tensor<?x?xf32>) -> (tensor<?xi64>)

    // CHECK: return %[[RESULT]]
    func.return %2 : tensor<?xi64>
  }

  // CHECK-LABEL: func @get_shape
  func.func @get_shape(%arg0 : tensor<*xi64>) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<*xi64>) -> tensor<?xi64>

                // to examine the phase it is associated to.
                if (execution.getPhase() != null) {
                    Map<Integer, List<MojoExecution>> phaseBindings = mappings.get(execution.getPhase());
                    if (phaseBindings != null) {
                        for (String goal : execution.getGoals()) {