__init__.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. # mypy: allow-untyped-defs
  2. from .modules import * # noqa: F403
  3. from .parameter import (
  4. Parameter as Parameter,
  5. UninitializedParameter as UninitializedParameter,
  6. UninitializedBuffer as UninitializedBuffer,
  7. )
  8. from .parallel import DataParallel as DataParallel
  9. from . import init
  10. from . import functional
  11. from . import utils
  12. from . import attention
  13. def factory_kwargs(kwargs):
  14. r"""Return a canonicalized dict of factory kwargs.
  15. Given kwargs, returns a canonicalized dict of factory kwargs that can be directly passed
  16. to factory functions like torch.empty, or errors if unrecognized kwargs are present.
  17. This function makes it simple to write code like this::
  18. class MyModule(nn.Module):
  19. def __init__(self, **kwargs):
  20. factory_kwargs = torch.nn.factory_kwargs(kwargs)
  21. self.weight = Parameter(torch.empty(10, **factory_kwargs))
  22. Why should you use this function instead of just passing `kwargs` along directly?
  23. 1. This function does error validation, so if there are unexpected kwargs we will
  24. immediately report an error, instead of deferring it to the factory call
  25. 2. This function supports a special `factory_kwargs` argument, which can be used to
  26. explicitly specify a kwarg to be used for factory functions, in the event one of the
  27. factory kwargs conflicts with an already existing argument in the signature (e.g.
  28. in the signature ``def f(dtype, **kwargs)``, you can specify ``dtype`` for factory
  29. functions, as distinct from the dtype argument, by saying
  30. ``f(dtype1, factory_kwargs={"dtype": dtype2})``)
  31. """
  32. if kwargs is None:
  33. return {}
  34. simple_keys = {"device", "dtype", "memory_format"}
  35. expected_keys = simple_keys | {"factory_kwargs"}
  36. if not kwargs.keys() <= expected_keys:
  37. raise TypeError(f"unexpected kwargs {kwargs.keys() - expected_keys}")
  38. # guarantee no input kwargs is untouched
  39. r = dict(kwargs.get("factory_kwargs", {}))
  40. for k in simple_keys:
  41. if k in kwargs:
  42. if k in r:
  43. raise TypeError(f"{k} specified twice, in **kwargs and in factory_kwargs")
  44. r[k] = kwargs[k]
  45. return r