图片解析应用
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

  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]: ...