if (metadata &&
      num_operands_expected != first_partitioned_input.getNumOperands()) {
    return first_partitioned_input->emitOpError()
           << "expects " << num_operands_expected << " operands but found "
           << first_partitioned_input.getNumOperands();
  }

  for (const auto& operand : replicated_input.getInputs().drop_front()) {
    auto partitioned_input =

QuantizeContext::QuantizeContext(func::FuncOp func, const DeviceTarget &spec)
    : func_(func), target_spec_(spec) {
  llvm::DenseMap<Value, int> value_to_state;
  func.walk([&](quantfork::QuantizeRegionOp op) {
    for (int i = 0, e = op.getNumOperands(); i != e; ++i) {
      states_manager_.InitializeOperandState(op, i, &value_to_state);
    }

    for (int res = 0, e = op.getNumResults(); res != e; ++res) {

// (2) Layout dependent arguments and results indices must be in
//     [0, getNumOperands/getNumResults) range.
LogicalResult VerifyLayoutSensitiveInterface(Operation* op);

// Verifies correctness of ops implementing FoldOperandsTransposeInterface (see
// definition in tf_op_base.td):
// (1) Layout dependent arguments and results indices must be in
//     [0, getNumOperands/getNumResults) range.

  std::vector<tensorflow::XlaExpression> expressions;
  std::vector<tensorflow::Tensor> tensors;
  std::vector<tensorflow::TensorValue> inputs;
  expressions.reserve(op_->getNumOperands());
  tensors.reserve(op_->getNumOperands());
  inputs.reserve(op_->getNumOperands());

  if (failed(
          PrepareKernelInputs(required_consts, expressions, tensors, inputs)))
    return failure();

  params_.inputs = inputs;

            mlir::func::ReturnOp::getOperationName());
  it++;  // tensor_cast
  it++;  // lstm
  EXPECT_EQ(it->getName().getStringRef(),
            mlir::TFL::LSTMOp::getOperationName());
  EXPECT_EQ(it->getNumOperands(), 24);
  EXPECT_EQ(it->getNumResults(), 1);
  // cifg = false, so input2input is not None.
  EXPECT_FALSE(mlir::isa<NoneType>(it->getOperand(1).getType()));
  // input layer norm is None

  const int num_operands = yield.getNumOperands();
  if (num_operands != result_count)
    return yield.emitOpError()
           << "has " << yield.getNumOperands()
           << " operand, but island returns " << result_count;
  for (int operand_idx : llvm::seq<int>(0, yield.getNumOperands())) {
    if (island.getResult(operand_idx).getType() !=

    if (!main) {
      op->emitError() << "MLIR Opaque Op expects a main() entry point\n";
      return;
    }
    if (main.getNumArguments() != op->getNumOperands()) {
      op->emitError() << "mismatch between MLIR Opaque Op number of operands ("
                      << op->getNumOperands()
                      << ") and main() entry point in the module ("
                      << main.getNumArguments() << " args)\n";
      return;
    }

        return success();
      }
    }
    return failure();
  }

  void rewrite(Operation* op, PatternRewriter& rewriter) const override {
    // Casts inputs of the operation.
    for (int i = 0; i < op->getNumOperands(); i++) {
      Value input = op->getOperand(i);
      if (getElementTypeOrSelf(input).isBF16()) {
        Value f32_cast = rewriter.create<TF::CastOp>(
            op->getLoc(),

  OpBuilder builder(&island.GetBody(), island.GetBody().begin());
  tf_executor::YieldOp yield = island.GetYield();
  if (yield.getNumOperands() == 0) {
    builder.create<TF::NoOp>(island.getLoc(), TypeRange{}, ValueRange{});
  } else if (yield.getNumOperands() == 1) {
    Value operand = yield.getOperand(0);
    auto identity = builder.create<TF::IdentityOp>(island.getLoc(),

    Operation* op, PatternRewriter& rewriter,
    const std::function<bool(ArrayRef<int64_t>, ArrayRef<int64_t>,
                             SmallVectorImpl<int64_t>&)>& get_broadcasted_shape)
    const {
  if (op->getNumOperands() != 2 || op->getResultTypes().size() != 1)
    return failure();

  // Check that the result shape is fully defined.
  auto result_type =
      mlir::dyn_cast_or_null<RankedTensorType>(op->getResultTypes().front());