};
  const SmallVector<Value> operands(dot_general_op->getOperands());
  const SmallVector<Value> results(dot_general_op->getResults());
  Operation* lifted_op =
      LiftAsFunctionCall(builder_, dot_general_op->getLoc(),
                         FunctionCallOpType::TFXlaCallModuleOp,
                         "composite_dot_general_fn", operands, results,
                         attributes)[0]
          .getDefiningOp();

    YieldOp yield_op = GetYield();
    // In order to correctly round-trip, we can only use this syntax when all
    // the locations are identical.
    if (controlled_op.getLoc() == getLoc() && yield_op.getLoc() == getLoc()) {
      p << " controls ";
      p.printGenericOp(&controlled_op);
      return;
    }
  }
  p << ' ';
  p.printRegion(getOperation()->getRegion(0));

    return failure();
  }

  std::vector<int64_t> lhs_shape = lhs_type.getShape();
  std::vector<int64_t> rhs_shape = rhs_type.getShape();

  Location loc = op.getLoc();

  // Ensure that input ranks are at least 2.
  const int lhs_dims = lhs_shape.size();
  const int rhs_dims = rhs_shape.size();
  if (lhs_dims < 2 || rhs_dims < 2) {
    // Both inputs must have rank >= 2

  auto fetches = llvm::to_vector(old_main_fetch.getFetches());
  fetches.emplace_back(value);

  auto builder = OpBuilder::atBlockTerminator(&graph_op.GetBody());
  builder.create<FetchOp>(old_main_fetch.getLoc(), std::move(fetches));
  old_main_fetch.erase();
}

void MergeSaveFunctionOpsToMain(func::FuncOp save_func_op,
                                func::FuncOp main_func_op) {

    // restrictive due to incompleteness of shape inference at this point.
    mlir::InFlightDiagnostic diag =
        mlir::emitWarning(op->getLoc(), "has static shape result");
    if (op->getContext()->shouldPrintOpOnDiagnostic()) {
      diag.attachNote(op->getLoc())
          .append("see current operation: ")
          .appendOp(*op, OpPrintingFlags().printGenericOpForm());
    }

    if (!erase[result.getResultNumber()]) {
      result_types.push_back(result.getType());
    }
  }
  OpBuilder builder(op);
  builder.setInsertionPointAfter(op);
  OperationState state(op->getLoc(), op->getName().getStringRef(),
                       op->getOperands(), result_types, op->getAttrs());
  for (int i = 0; i < op->getNumRegions(); ++i) {
    state.addRegion();
  }

      };

      // Insert custom aggregation op between operand and operator.
      rewriter.setInsertionPointAfterValue(value);
      Operation *aggregator_op = rewriter.create<TF::CustomAggregatorOp>(
          op->getLoc(), output_types, value, attributes);

      Value aggregator_op_result = aggregator_op->getOpResult(0);
      value.replaceAllUsesExcept(aggregator_op_result, aggregator_op);
    }

    return success();
  }

            return !wrapped_ops.count(user.getOwner());
          }))
        result_types.push_back(result.getType());
    }
  }

  IslandOp new_island = OpBuilder(island).create<IslandOp>(
      island.getLoc(), result_types,
      /*control=*/ControlType::get(island.getContext()),
      /*controlInputs=*/island.getOperands());
  new_island.getBody().push_back(new Block);

  // resolved.
  Type type;
  TF_RETURN_IF_ERROR(ConvertDataTypeToTensor(dtype, builder_, &type));
  *handle =
      new MlirTensor(func_.getBody().front().addArgument(type, func_.getLoc()));
  return absl::OkStatus();
}

Status MlirAbstractOp::AddInput(AbstractTensorHandle* input) {
  if (current_ods_input_ >= op_def_->input_arg_size())
    return InvalidArgument(

      new_operands.push_back(while_op.getOperand(i));
      new_result_types.push_back(while_op.getResultTypes()[i]);
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {