mlir::TF::_XlaRecvAtHostV2Op>(defining_op)) {
            external_values.insert(v);
          }
          continue;
        }
        auto block_arg = mlir::cast<BlockArgument>(v);
        if (block_arg.getParentRegion() == op->getParentRegion())
          external_values.insert(v);
      }
      return WalkResult::advance();
    });
  }
  return external_values;
}

                     parent.getDefiningOp<TF::TensorListSetItemOp>()) {
        parent = set_item.getInputHandle();
      } else {
        break;
      }
    }
    if (auto block_arg = parent.dyn_cast<mlir::BlockArgument>()) {
      return block_arg.getArgNumber();
    }
    // Returns -1 if we don't find which this result maps to.
    return -1;
  };

  llvm::SmallVector<Value, 4> returns;

  //
  // Note: This does not descend through nested SCPs.
  auto ComesFromBlockArgNumber = [](Value val) -> int {
    while (true) {
      if (auto block_arg = llvm::dyn_cast<BlockArgument>(val)) {
        return block_arg.getArgNumber();
      }
      if (auto identity_op =
              llvm::dyn_cast<TF::IdentityOp>(val.getDefiningOp())) {
        val = identity_op.getOperand();
      } else {