pytorch / pytorch/pytorch

Add a lazy import hook for autoloading?

Open
#180,712 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged feature module: backend module: ux triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🚀 The feature, motivation and pitch

This probably lands somewhere between #180343 and #180345, not sure if it is justified in having its own issue, but in TorchTPU we have set up autoloading of our `tpu` module, but we only load it if a TPU is present, since otherwise it won't work anyway. The problem is that this leads to a pretty bad error message, since what ends up happening is that when you do `torch.device("tpu")` it gives this error:

```
RuntimeError: Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip,
ve, fpga, maia, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone
device type at start of device string: tpu
```

This doesn't make it clear that the problem is that `torch_tpu` is deliberately choosing *not* to load the backend because the TPU device is not available. We could throw an error instead of choosing not to load, but that means any time you install a backend on machine where it can't load, you would be unable to import torch *at all*.

Ideally, we would be able to register a hook that would be called the first time someone tries to actually *use* our device or module, so that we have a stronger signal that an error should be thrown than "User has our package installed" (which may just be because it ended up in a `requirements.txt` that doesn't sufficiently distinguish between environments with and without a TPU).

This would have also helped us with an earlier problem we had where we wanted to defer full initialization until first use. We've since basically moved that logic to the C++ layer, but it would have been very convenient if instead of registering our backend, we could register lazy initialization hooks.

### Alternatives

With unlimited backends like #180345, I would be a lot more comfortable writing a generic backend that throws an error whenever you try to load something from it to explain that no TPU is available.

Alternatively, if there were a simple API to explicitly register "This name is valid but an error should be thrown when you try to use it", that would solve our immediate use case as well.

### Additional context

_No response_

cc @bdhirsh

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.