parameter.pyi 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. from typing_extensions import TypeIs
  2. from torch import device, dtype, Tensor
  3. class Parameter(Tensor):
  4. def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
  5. def is_lazy(
  6. param: Tensor,
  7. ) -> TypeIs[UninitializedParameter | UninitializedBuffer]: ...
  8. class UninitializedParameter(Tensor):
  9. def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
  10. def materialize(
  11. self,
  12. shape: tuple[int, ...],
  13. device: device | None = None,
  14. dtype: dtype | None = None,
  15. ) -> None: ...
  16. class Buffer(Tensor):
  17. persistent: bool
  18. def __init__(
  19. self,
  20. data: Tensor = ...,
  21. requires_grad: bool = ...,
  22. persistent: bool = ...,
  23. ) -> None: ...
  24. class UninitializedBuffer(Tensor):
  25. persistent: bool
  26. def __init__(
  27. self,
  28. data: Tensor = ...,
  29. requires_grad: bool = ...,
  30. persistent: bool = ...,
  31. ) -> None: ...
  32. def materialize(
  33. self,
  34. shape: tuple[int, ...],
  35. device: device | None = None,
  36. dtype: dtype | None = None,
  37. ) -> None: ...