p.printOperands(op->getOperands());
  p.printOptionalAttrDict(op->getAttrs());
  // Now we can output only one type for all operands and the result.
  p << " : " << resultType;
}

Operation* getDefiningBroadcastArgsOp(Value operand) {
  auto* defining_op = operand.getDefiningOp();
  if (!llvm::dyn_cast_or_null<TF::BroadcastToOp>(defining_op) &&
      !llvm::dyn_cast_or_null<TFL::BroadcastToOp>(defining_op)) {
    return nullptr;