Add a lazy import hook for autoloading?
- 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
Assessment
This issue has not been assessed yet.