pytorch - 💡(How to fix) Fix Add a lazy import hook for autoloading? [1 participants]

🚀 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