[](mlir::OperationState& state, std::vector<mlir::Type> tys) {
             state.addTypes(mlir::ArrayRef<mlir::Type>(tys));
           })
      .def("addOperands",
           [](mlir::OperationState& os, std::vector<mlir::Value> ops) {
             os.addOperands(mlir::ArrayRef<mlir::Value>(ops));
           })
      .def("addRegion", py::overload_cast<>(&mlir::OperationState::addRegion),
           py::return_value_policy::reference);

    // Build the newly fused operation to replace the batch norm
    OperationState state(batch_norm.getLoc(),
                         _FusedBatchNormExOp::getOperationName());
    state.addOperands(batch_norm.getOperands());
    if (side_input) state.operands.push_back(side_input);
    state.addTypes(batch_norm.getResultTypes());
    state.addAttributes(batch_norm->getAttrs());

    def getRegion(self, arg0: int) -> Region: ...
    def getResult(self, arg0: int) -> OpResult: ...

class OperationState:
    def __init__(self, arg0: Location, arg1: str) -> None: ...
    def addOperands(self, arg0: list[Value]) -> None: ...
    def addRegion(self) -> Region: ...
    def addTypes(self, arg0: list[Type]) -> None: ...

class RankedTensorType(Type):
    def __init__(self, *args, **kwargs) -> None: ...

    auto attr =
        ReadAttr(flex_buffer_map, &builder, custom_op.getCustomCode().str());
    OperationState op_state(custom_op.getLoc(),
                            custom_op.getCustomCode().str());
    op_state.addOperands(custom_op.getOperands());
    llvm::SmallVector<mlir::Type, 4> output_tys;
    for (int i = 0; i < custom_op.getNumResults(); i++) {
      output_tys.push_back(custom_op.getType(i));
    }

    const std::string tf_op_full_name = llvm::Twine("tf.", tf_op_name).str();

    // Create the TF op
    OperationState op_state(op.getLoc(), tf_op_full_name);
    op_state.addOperands(op.getOperands());
    op_state.addTypes(op.getResultTypes());

    SmallVector<NamedAttribute, 2> attrs;
    std::string parsed_op_name;
    tensorflow::NodeDef node_def;

  mlir::OpBuilder op_builder(func.getBody());

  // Create the TF op
  const std::string tf_op_full_name = absl::StrCat("tf.", node_def.op());
  mlir::OperationState op_state(loc, tf_op_full_name);
  op_state.addOperands(func.getArguments());
  op_state.addTypes(output_tys);
  op_state.addAttributes(attrs);
  mlir::Operation* tf_op = op_builder.create(op_state);
  op_builder.create<mlir::func::ReturnOp>(loc, tf_op->getResults());

  OperationState op_state(loc, op_name);

  for (auto input_num : op.inputs) {
    if (input_num == -1) {
      assert(optional_arg_marker != nullptr);
      op_state.addOperands({optional_arg_marker});
    } else {
      op_state.addOperands({vals_map.at(input_num)});
    }
  }

  for (auto output_num : op.outputs) {
    auto& tensor = *tensors.at(output_num);

      DCHECK(succeeded(
          VerifyCompatibleTypes(input.getType(), replicated_input.second)));
      state->addOperands(input);
    }
    block.addArgument(replicated_input.second, state->location);
  }

  for (auto packed_input : packed_inputs) {
    state->addOperands(packed_input);
    block.addArgument(packed_input.getType(), state->location);
  }

      // op for some other operand. Update defining_ops.
      if (defining_ops.insert(new_control_input.getDefiningOp()).second) {
        operands.push_back(new_control_input);
      }
    }
    state.addOperands(operands);
    Operation* new_op = builder.create(state);
    item.replaceAllUsesWith(new_op);
    new_op->setAttrs(item.getAttrDictionary());
    item.erase();
  }
  AddStatefulControlFlowDependencies(graph_op);

  [{
    auto resultType =
      OpTrait::util::getBroadcastedType(lhs.getType(), rhs.getType());
    if (!resultType)
      mlir::emitError($_state.location, "non-broadcastable operands");
    $_state.addOperands({lhs, rhs});
    $_state.types.push_back(resultType);
  }]>;

def TFL_FusedBroadcastableBinaryBuilder :
  OpBuilder<(ins "Value":$lhs, "Value":$rhs,
    "StringAttr":$fusedActivationFunction),
  [{