backend_registration.py 1.3 KB

12345678910111213141516171819202122232425262728293031
  1. from torch._C import _rename_privateuse1_backend
  2. def rename_privateuse1_backend(backend_name: str) -> None:
  3. r"""
  4. rename_privateuse1_backend(backend_name) -> None
  5. This is a registration API for external backends that would like to register their
  6. own device and C++ kernels out of tree.
  7. The steps are:
  8. (1) (In C++) implement kernels for various torch operations, and register them
  9. to the PrivateUse1 dispatch key.
  10. (2) (In python) call torch.register_privateuse1_backend("foo")
  11. You can now use "foo" as an ordinary device string in python.
  12. Note: this API can only be called once per process. Attempting to change
  13. the external backend after it's already been set will result in an error.
  14. For more details, see https://pytorch.org/tutorials/advanced/extend_dispatcher.html#get-a-dispatch-key-for-your-backend
  15. For an existing example, see https://github.com/bdhirsh/pytorch_open_registration_example
  16. Example::
  17. >>> # xdoctest: +SKIP("failing")
  18. >>> torch.register_privateuse1_backend("foo")
  19. # This will work, assuming that you've implemented the right C++ kernels
  20. # to implement torch.ones.
  21. >>> a = torch.ones(2, device="foo")
  22. """
  23. return _rename_privateuse1_backend(backend_name)