m2m模型翻译
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

142 lines
3.6 KiB

6 months ago
  1. import sys
  2. from typing import List, TypeVar, Optional, Any, overload, Union, Tuple, Sequence
  3. from numpy import (
  4. ndarray,
  5. dtype,
  6. bool_,
  7. unsignedinteger,
  8. signedinteger,
  9. floating,
  10. complexfloating,
  11. number,
  12. _OrderKACF,
  13. )
  14. from numpy.typing import (
  15. _ArrayLikeBool_co,
  16. _ArrayLikeUInt_co,
  17. _ArrayLikeInt_co,
  18. _ArrayLikeFloat_co,
  19. _ArrayLikeComplex_co,
  20. _DTypeLikeBool,
  21. _DTypeLikeUInt,
  22. _DTypeLikeInt,
  23. _DTypeLikeFloat,
  24. _DTypeLikeComplex,
  25. _DTypeLikeComplex_co,
  26. )
  27. if sys.version_info >= (3, 8):
  28. from typing import Literal
  29. else:
  30. from typing_extensions import Literal
  31. _ArrayType = TypeVar(
  32. "_ArrayType",
  33. bound=ndarray[Any, dtype[Union[bool_, number[Any]]]],
  34. )
  35. _OptimizeKind = Union[
  36. None, bool, Literal["greedy", "optimal"], Sequence[Any]
  37. ]
  38. _CastingSafe = Literal["no", "equiv", "safe", "same_kind"]
  39. _CastingUnsafe = Literal["unsafe"]
  40. __all__: List[str]
  41. # TODO: Properly handle the `casting`-based combinatorics
  42. # TODO: We need to evaluate the content `__subscripts` in order
  43. # to identify whether or an array or scalar is returned. At a cursory
  44. # glance this seems like something that can quite easilly be done with
  45. # a mypy plugin.
  46. # Something like `is_scalar = bool(__subscripts.partition("->")[-1])`
  47. @overload
  48. def einsum(
  49. __subscripts: str,
  50. *operands: _ArrayLikeBool_co,
  51. out: None = ...,
  52. dtype: Optional[_DTypeLikeBool] = ...,
  53. order: _OrderKACF = ...,
  54. casting: _CastingSafe = ...,
  55. optimize: _OptimizeKind = ...,
  56. ) -> Any: ...
  57. @overload
  58. def einsum(
  59. __subscripts: str,
  60. *operands: _ArrayLikeUInt_co,
  61. out: None = ...,
  62. dtype: Optional[_DTypeLikeUInt] = ...,
  63. order: _OrderKACF = ...,
  64. casting: _CastingSafe = ...,
  65. optimize: _OptimizeKind = ...,
  66. ) -> Any: ...
  67. @overload
  68. def einsum(
  69. __subscripts: str,
  70. *operands: _ArrayLikeInt_co,
  71. out: None = ...,
  72. dtype: Optional[_DTypeLikeInt] = ...,
  73. order: _OrderKACF = ...,
  74. casting: _CastingSafe = ...,
  75. optimize: _OptimizeKind = ...,
  76. ) -> Any: ...
  77. @overload
  78. def einsum(
  79. __subscripts: str,
  80. *operands: _ArrayLikeFloat_co,
  81. out: None = ...,
  82. dtype: Optional[_DTypeLikeFloat] = ...,
  83. order: _OrderKACF = ...,
  84. casting: _CastingSafe = ...,
  85. optimize: _OptimizeKind = ...,
  86. ) -> Any: ...
  87. @overload
  88. def einsum(
  89. __subscripts: str,
  90. *operands: _ArrayLikeComplex_co,
  91. out: None = ...,
  92. dtype: Optional[_DTypeLikeComplex] = ...,
  93. order: _OrderKACF = ...,
  94. casting: _CastingSafe = ...,
  95. optimize: _OptimizeKind = ...,
  96. ) -> Any: ...
  97. @overload
  98. def einsum(
  99. __subscripts: str,
  100. *operands: Any,
  101. casting: _CastingUnsafe,
  102. dtype: Optional[_DTypeLikeComplex_co] = ...,
  103. out: None = ...,
  104. order: _OrderKACF = ...,
  105. optimize: _OptimizeKind = ...,
  106. ) -> Any: ...
  107. @overload
  108. def einsum(
  109. __subscripts: str,
  110. *operands: _ArrayLikeComplex_co,
  111. out: _ArrayType,
  112. dtype: Optional[_DTypeLikeComplex_co] = ...,
  113. order: _OrderKACF = ...,
  114. casting: _CastingSafe = ...,
  115. optimize: _OptimizeKind = ...,
  116. ) -> _ArrayType: ...
  117. @overload
  118. def einsum(
  119. __subscripts: str,
  120. *operands: Any,
  121. out: _ArrayType,
  122. casting: _CastingUnsafe,
  123. dtype: Optional[_DTypeLikeComplex_co] = ...,
  124. order: _OrderKACF = ...,
  125. optimize: _OptimizeKind = ...,
  126. ) -> _ArrayType: ...
  127. # NOTE: `einsum_call` is a hidden kwarg unavailable for public use.
  128. # It is therefore excluded from the signatures below.
  129. # NOTE: In practice the list consists of a `str` (first element)
  130. # and a variable number of integer tuples.
  131. def einsum_path(
  132. __subscripts: str,
  133. *operands: _ArrayLikeComplex_co,
  134. optimize: _OptimizeKind = ...,
  135. ) -> Tuple[List[Any], str]: ...