nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (309) hide show
  1. nvmath/__init__.pxd +0 -0
  2. nvmath/__init__.py +45 -0
  3. nvmath/_internal/__init__.py +0 -0
  4. nvmath/_internal/attribute_ifc_factory.py +330 -0
  5. nvmath/_internal/layout.py +70 -0
  6. nvmath/_internal/templates.py +383 -0
  7. nvmath/_utils.py +150 -0
  8. nvmath/bindings/__init__.py +51 -0
  9. nvmath/bindings/_internal/__init__.pxd +0 -0
  10. nvmath/bindings/_internal/__init__.py +0 -0
  11. nvmath/bindings/_internal/cublas.cp314-win_amd64.pyd +0 -0
  12. nvmath/bindings/_internal/cublas.pxd +530 -0
  13. nvmath/bindings/_internal/cublasLt.cp314-win_amd64.pyd +0 -0
  14. nvmath/bindings/_internal/cublasLt.pxd +59 -0
  15. nvmath/bindings/_internal/cublasMp.pxd +52 -0
  16. nvmath/bindings/_internal/cudss.cp314-win_amd64.pyd +0 -0
  17. nvmath/bindings/_internal/cudss.pxd +47 -0
  18. nvmath/bindings/_internal/cufft.cp314-win_amd64.pyd +0 -0
  19. nvmath/bindings/_internal/cufft.pxd +70 -0
  20. nvmath/bindings/_internal/cufftMp.pxd +77 -0
  21. nvmath/bindings/_internal/curand.cp314-win_amd64.pyd +0 -0
  22. nvmath/bindings/_internal/curand.pxd +42 -0
  23. nvmath/bindings/_internal/cusolver.cp314-win_amd64.pyd +0 -0
  24. nvmath/bindings/_internal/cusolver.pxd +15 -0
  25. nvmath/bindings/_internal/cusolverDn.cp314-win_amd64.pyd +0 -0
  26. nvmath/bindings/_internal/cusolverDn.pxd +406 -0
  27. nvmath/bindings/_internal/cusparse.cp314-win_amd64.pyd +0 -0
  28. nvmath/bindings/_internal/cusparse.pxd +277 -0
  29. nvmath/bindings/_internal/cusparseLt.cp314-win_amd64.pyd +0 -0
  30. nvmath/bindings/_internal/cusparseLt.pxd +48 -0
  31. nvmath/bindings/_internal/cutensor.cp314-win_amd64.pyd +0 -0
  32. nvmath/bindings/_internal/cutensor.pxd +58 -0
  33. nvmath/bindings/_internal/mathdx.cp314-win_amd64.pyd +0 -0
  34. nvmath/bindings/_internal/mathdx.pxd +116 -0
  35. nvmath/bindings/_internal/nvshmem.pxd +29 -0
  36. nvmath/bindings/_internal/utils.cp314-win_amd64.pyd +0 -0
  37. nvmath/bindings/_internal/utils.pxd +179 -0
  38. nvmath/bindings/_internal/utils.pyi +10 -0
  39. nvmath/bindings/cublas.cp314-win_amd64.pyd +0 -0
  40. nvmath/bindings/cublas.pxd +558 -0
  41. nvmath/bindings/cublas.pyi +746 -0
  42. nvmath/bindings/cublasLt.cp314-win_amd64.pyd +0 -0
  43. nvmath/bindings/cublasLt.pxd +109 -0
  44. nvmath/bindings/cublasLt.pyi +1337 -0
  45. nvmath/bindings/cublasMp.pxd +85 -0
  46. nvmath/bindings/cublasMp.pyi +219 -0
  47. nvmath/bindings/cudss.cp314-win_amd64.pyd +0 -0
  48. nvmath/bindings/cudss.pxd +86 -0
  49. nvmath/bindings/cudss.pyi +282 -0
  50. nvmath/bindings/cufft.cp314-win_amd64.pyd +0 -0
  51. nvmath/bindings/cufft.pxd +118 -0
  52. nvmath/bindings/cufft.pyi +241 -0
  53. nvmath/bindings/cufftMp.pxd +124 -0
  54. nvmath/bindings/cufftMp.pyi +260 -0
  55. nvmath/bindings/curand.cp314-win_amd64.pyd +0 -0
  56. nvmath/bindings/curand.pxd +71 -0
  57. nvmath/bindings/curand.pyi +159 -0
  58. nvmath/bindings/cusolver.cp314-win_amd64.pyd +0 -0
  59. nvmath/bindings/cusolver.pxd +62 -0
  60. nvmath/bindings/cusolver.pyi +242 -0
  61. nvmath/bindings/cusolverDn.cp314-win_amd64.pyd +0 -0
  62. nvmath/bindings/cusolverDn.pxd +430 -0
  63. nvmath/bindings/cusolverDn.pyi +416 -0
  64. nvmath/bindings/cusparse.cp314-win_amd64.pyd +0 -0
  65. nvmath/bindings/cusparse.pxd +338 -0
  66. nvmath/bindings/cusparse.pyi +659 -0
  67. nvmath/bindings/cusparseLt.cp314-win_amd64.pyd +0 -0
  68. nvmath/bindings/cusparseLt.pxd +99 -0
  69. nvmath/bindings/cusparseLt.pyi +198 -0
  70. nvmath/bindings/cutensor.cp314-win_amd64.pyd +0 -0
  71. nvmath/bindings/cutensor.pxd +98 -0
  72. nvmath/bindings/cutensor.pyi +294 -0
  73. nvmath/bindings/cycublas.cp314-win_amd64.pyd +0 -0
  74. nvmath/bindings/cycublas.pxd +684 -0
  75. nvmath/bindings/cycublasLt.cp314-win_amd64.pyd +0 -0
  76. nvmath/bindings/cycublasLt.pxd +1055 -0
  77. nvmath/bindings/cycublasMp.pxd +183 -0
  78. nvmath/bindings/cycudss.cp314-win_amd64.pyd +0 -0
  79. nvmath/bindings/cycudss.pxd +224 -0
  80. nvmath/bindings/cycufft.cp314-win_amd64.pyd +0 -0
  81. nvmath/bindings/cycufft.pxd +326 -0
  82. nvmath/bindings/cycufftMp.pxd +334 -0
  83. nvmath/bindings/cycurand.cp314-win_amd64.pyd +0 -0
  84. nvmath/bindings/cycurand.pxd +146 -0
  85. nvmath/bindings/cycusolver.cp314-win_amd64.pyd +0 -0
  86. nvmath/bindings/cycusolver.pxd +154 -0
  87. nvmath/bindings/cycusolverDn.cp314-win_amd64.pyd +0 -0
  88. nvmath/bindings/cycusolverDn.pxd +446 -0
  89. nvmath/bindings/cycusparse.cp314-win_amd64.pyd +0 -0
  90. nvmath/bindings/cycusparse.pxd +470 -0
  91. nvmath/bindings/cycusparseLt.cp314-win_amd64.pyd +0 -0
  92. nvmath/bindings/cycusparseLt.pxd +150 -0
  93. nvmath/bindings/cycutensor.cp314-win_amd64.pyd +0 -0
  94. nvmath/bindings/cycutensor.pxd +192 -0
  95. nvmath/bindings/cymathdx.cp314-win_amd64.pyd +0 -0
  96. nvmath/bindings/cymathdx.pxd +552 -0
  97. nvmath/bindings/cynvshmem.pxd +126 -0
  98. nvmath/bindings/mathdx.cp314-win_amd64.pyd +0 -0
  99. nvmath/bindings/mathdx.pxd +182 -0
  100. nvmath/bindings/mathdx.pyi +1244 -0
  101. nvmath/bindings/nvpl/__init__.pxd +0 -0
  102. nvmath/bindings/nvpl/__init__.py +13 -0
  103. nvmath/bindings/nvpl/_internal/__init__.pxd +0 -0
  104. nvmath/bindings/nvpl/_internal/__init__.py +0 -0
  105. nvmath/bindings/nvpl/_internal/blas.cp314-win_amd64.pyd +0 -0
  106. nvmath/bindings/nvpl/_internal/blas.pxd +125 -0
  107. nvmath/bindings/nvpl/_internal/fft.cp314-win_amd64.pyd +0 -0
  108. nvmath/bindings/nvpl/_internal/fft.pxd +36 -0
  109. nvmath/bindings/nvpl/blas.cp314-win_amd64.pyd +0 -0
  110. nvmath/bindings/nvpl/blas.pxd +85 -0
  111. nvmath/bindings/nvpl/blas.pyi +122 -0
  112. nvmath/bindings/nvpl/cyblas.cp314-win_amd64.pyd +0 -0
  113. nvmath/bindings/nvpl/cyblas.pxd +166 -0
  114. nvmath/bindings/nvpl/cyfft.cp314-win_amd64.pyd +0 -0
  115. nvmath/bindings/nvpl/cyfft.pxd +92 -0
  116. nvmath/bindings/nvpl/fft.cp314-win_amd64.pyd +0 -0
  117. nvmath/bindings/nvpl/fft.pxd +100 -0
  118. nvmath/bindings/nvpl/fft.pyi +100 -0
  119. nvmath/bindings/nvshmem.pxd +54 -0
  120. nvmath/bindings/nvshmem.pyi +179 -0
  121. nvmath/device/__init__.py +21 -0
  122. nvmath/device/_deprecated.py +33 -0
  123. nvmath/device/common.py +313 -0
  124. nvmath/device/common_backend.py +131 -0
  125. nvmath/device/common_cuda.py +201 -0
  126. nvmath/device/common_numba.py +310 -0
  127. nvmath/device/common_opaque_tensor.py +201 -0
  128. nvmath/device/cublasdx.py +1805 -0
  129. nvmath/device/cublasdx_backend.py +807 -0
  130. nvmath/device/cublasdx_numba.py +1612 -0
  131. nvmath/device/cufftdx.py +510 -0
  132. nvmath/device/cufftdx_backend.py +196 -0
  133. nvmath/device/cufftdx_numba.py +196 -0
  134. nvmath/device/curand_kernel.py +9147 -0
  135. nvmath/device/cusolverdx.py +2690 -0
  136. nvmath/device/cusolverdx_backend.py +440 -0
  137. nvmath/device/cusolverdx_numba.py +1624 -0
  138. nvmath/device/llvm_array.py +29 -0
  139. nvmath/device/random.py +445 -0
  140. nvmath/device/random_helpers.py +23 -0
  141. nvmath/device/random_states.py +187 -0
  142. nvmath/device/types.py +147 -0
  143. nvmath/device/vector_types_numba.py +203 -0
  144. nvmath/distributed/__init__.py +205 -0
  145. nvmath/distributed/_internal/__init__.py +0 -0
  146. nvmath/distributed/_internal/nvshmem.py +302 -0
  147. nvmath/distributed/_internal/tensor_ifc.py +67 -0
  148. nvmath/distributed/_internal/tensor_ifc_cupy.py +62 -0
  149. nvmath/distributed/_internal/tensor_ifc_host_device.py +165 -0
  150. nvmath/distributed/_internal/tensor_ifc_numpy.py +41 -0
  151. nvmath/distributed/_internal/tensor_ifc_torch.py +141 -0
  152. nvmath/distributed/_internal/tensor_wrapper.py +78 -0
  153. nvmath/distributed/_utils.py +167 -0
  154. nvmath/distributed/distribution.py +770 -0
  155. nvmath/distributed/fft/__init__.py +7 -0
  156. nvmath/distributed/fft/_configuration.py +79 -0
  157. nvmath/distributed/fft/fft.py +2801 -0
  158. nvmath/distributed/linalg/__init__.py +12 -0
  159. nvmath/distributed/linalg/_internal/__init__.py +3 -0
  160. nvmath/distributed/linalg/_internal/epilog_protocol.py +496 -0
  161. nvmath/distributed/linalg/_internal/matmul_desc_ifc.py +28 -0
  162. nvmath/distributed/linalg/advanced/__init__.py +8 -0
  163. nvmath/distributed/linalg/advanced/_configuration.py +166 -0
  164. nvmath/distributed/linalg/advanced/matmulmod.py +2908 -0
  165. nvmath/distributed/process_group.py +408 -0
  166. nvmath/distributed/reshape/__init__.py +6 -0
  167. nvmath/distributed/reshape/_configuration.py +39 -0
  168. nvmath/distributed/reshape/reshape.py +1256 -0
  169. nvmath/fft/__init__.py +7 -0
  170. nvmath/fft/_configuration.py +208 -0
  171. nvmath/fft/_exec_utils.py +109 -0
  172. nvmath/fft/_helpers.py +237 -0
  173. nvmath/fft/fft.py +2903 -0
  174. nvmath/internal/__init__.pxd +3 -0
  175. nvmath/internal/__init__.py +10 -0
  176. nvmath/internal/_bindings.cp314-win_amd64.pyd +0 -0
  177. nvmath/internal/_bindings.pxd +12 -0
  178. nvmath/internal/_device_utils.py +54 -0
  179. nvmath/internal/enum_utils.py +142 -0
  180. nvmath/internal/formatters.py +87 -0
  181. nvmath/internal/mem_limit.py +51 -0
  182. nvmath/internal/memory.cp314-win_amd64.pyd +0 -0
  183. nvmath/internal/memory.pxd +10 -0
  184. nvmath/internal/memory.pyi +46 -0
  185. nvmath/internal/ndbuffer/__init__.pxd +3 -0
  186. nvmath/internal/ndbuffer/__init__.py +7 -0
  187. nvmath/internal/ndbuffer/copy_kernel/args.h +34 -0
  188. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/array_view.h +52 -0
  189. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/elementwise.h +68 -0
  190. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/grid_indexer.h +69 -0
  191. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/transposed.h +240 -0
  192. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/type_utils.h +39 -0
  193. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/utils.h +132 -0
  194. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/vec.h +159 -0
  195. nvmath/internal/ndbuffer/copy_kernel/elementwise.h +53 -0
  196. nvmath/internal/ndbuffer/copy_kernel/transposed.h +58 -0
  197. nvmath/internal/ndbuffer/copy_kernel.cp314-win_amd64.pyd +0 -0
  198. nvmath/internal/ndbuffer/copy_kernel.pxd +9 -0
  199. nvmath/internal/ndbuffer/copy_kernel.pyi +9 -0
  200. nvmath/internal/ndbuffer/data_layout.cp314-win_amd64.pyd +0 -0
  201. nvmath/internal/ndbuffer/data_layout.pxd +72 -0
  202. nvmath/internal/ndbuffer/data_layout.pyi +23 -0
  203. nvmath/internal/ndbuffer/jit.cp314-win_amd64.pyd +0 -0
  204. nvmath/internal/ndbuffer/jit.pxd +7 -0
  205. nvmath/internal/ndbuffer/jit.pyi +6 -0
  206. nvmath/internal/ndbuffer/ndbuffer.cp314-win_amd64.pyd +0 -0
  207. nvmath/internal/ndbuffer/ndbuffer.pxd +36 -0
  208. nvmath/internal/ndbuffer/ndbuffer.pyi +41 -0
  209. nvmath/internal/ndbuffer/package_utils.cp314-win_amd64.pyd +0 -0
  210. nvmath/internal/ndbuffer/package_utils.pxd +11 -0
  211. nvmath/internal/ndbuffer/package_utils.pyi +12 -0
  212. nvmath/internal/package_ifc.py +147 -0
  213. nvmath/internal/package_ifc_cuda.py +49 -0
  214. nvmath/internal/package_ifc_cupy.py +49 -0
  215. nvmath/internal/package_ifc_torch.py +31 -0
  216. nvmath/internal/package_wrapper.py +14 -0
  217. nvmath/internal/tensor_ifc.py +194 -0
  218. nvmath/internal/tensor_ifc_cupy.py +241 -0
  219. nvmath/internal/tensor_ifc_ndbuffer.py +153 -0
  220. nvmath/internal/tensor_ifc_numpy.py +183 -0
  221. nvmath/internal/tensor_ifc_torch.py +176 -0
  222. nvmath/internal/tensor_wrapper.py +157 -0
  223. nvmath/internal/typemaps.py +162 -0
  224. nvmath/internal/utils.py +771 -0
  225. nvmath/linalg/__init__.py +46 -0
  226. nvmath/linalg/_internal/__init__.py +3 -0
  227. nvmath/linalg/_internal/algo_cap_ifc.py +82 -0
  228. nvmath/linalg/_internal/algo_config_ifc.py +43 -0
  229. nvmath/linalg/_internal/batch.py +217 -0
  230. nvmath/linalg/_internal/enum_to_tuples.py +64 -0
  231. nvmath/linalg/_internal/epilog_protocol.py +766 -0
  232. nvmath/linalg/_internal/layout.py +624 -0
  233. nvmath/linalg/_internal/matmul_desc_ifc.py +28 -0
  234. nvmath/linalg/_internal/matmul_pref_ifc.py +27 -0
  235. nvmath/linalg/_internal/matrix_layout_ifc.py +26 -0
  236. nvmath/linalg/_internal/typemaps.py +144 -0
  237. nvmath/linalg/_internal/utils.py +134 -0
  238. nvmath/linalg/advanced/__init__.py +8 -0
  239. nvmath/linalg/advanced/_algorithmmod.py +149 -0
  240. nvmath/linalg/advanced/_configuration.py +349 -0
  241. nvmath/linalg/advanced/helpers/__init__.py +5 -0
  242. nvmath/linalg/advanced/helpers/matmul.py +1348 -0
  243. nvmath/linalg/advanced/matmulmod.py +3937 -0
  244. nvmath/linalg/generic/__init__.py +43 -0
  245. nvmath/linalg/generic/_configuration/__init__.py +37 -0
  246. nvmath/linalg/generic/_configuration/layout.py +263 -0
  247. nvmath/linalg/generic/_configuration/match.py +609 -0
  248. nvmath/linalg/generic/_configuration/qualifiers.py +494 -0
  249. nvmath/linalg/generic/_configuration/wrap.py +217 -0
  250. nvmath/linalg/generic/_dtype.py +15 -0
  251. nvmath/linalg/generic/matmulmod.py +1060 -0
  252. nvmath/memory.py +282 -0
  253. nvmath/sparse/__init__.py +38 -0
  254. nvmath/sparse/_internal/__init__.py +21 -0
  255. nvmath/sparse/_internal/common_utils.py +144 -0
  256. nvmath/sparse/_internal/cudss_config_ifc.py +741 -0
  257. nvmath/sparse/_internal/cudss_data_ifc.py +399 -0
  258. nvmath/sparse/_internal/cudss_utils.py +445 -0
  259. nvmath/sparse/_internal/cusparse_utils.py +382 -0
  260. nvmath/sparse/_internal/sparse_bsc_ifc.py +300 -0
  261. nvmath/sparse/_internal/sparse_bsr_ifc.py +303 -0
  262. nvmath/sparse/_internal/sparse_coo_ifc.py +254 -0
  263. nvmath/sparse/_internal/sparse_csc_ifc.py +266 -0
  264. nvmath/sparse/_internal/sparse_csr_ifc.py +266 -0
  265. nvmath/sparse/_internal/sparse_dia_ifc.py +239 -0
  266. nvmath/sparse/_internal/sparse_format_helpers.py +601 -0
  267. nvmath/sparse/_internal/sparse_tensor_ifc.py +121 -0
  268. nvmath/sparse/_internal/sparse_ust_ifc.py +146 -0
  269. nvmath/sparse/_internal/utils.py +56 -0
  270. nvmath/sparse/advanced/__init__.py +7 -0
  271. nvmath/sparse/advanced/_configuration.py +230 -0
  272. nvmath/sparse/advanced/direct_solver.py +1884 -0
  273. nvmath/sparse/generic/__init__.py +7 -0
  274. nvmath/sparse/generic/_configuration.py +129 -0
  275. nvmath/sparse/generic/_helpers.py +137 -0
  276. nvmath/sparse/generic/_thunks.py +21 -0
  277. nvmath/sparse/generic/matmulmod.py +2426 -0
  278. nvmath/sparse/ust/__init__.py +7 -0
  279. nvmath/sparse/ust/_converters.py +378 -0
  280. nvmath/sparse/ust/_drawer.py +552 -0
  281. nvmath/sparse/ust/_emitter.py +553 -0
  282. nvmath/sparse/ust/_jit.py +212 -0
  283. nvmath/sparse/ust/_kernel.py +346 -0
  284. nvmath/sparse/ust/_semiring.py +71 -0
  285. nvmath/sparse/ust/_utils.py +102 -0
  286. nvmath/sparse/ust/interfaces/__init__.py +0 -0
  287. nvmath/sparse/ust/interfaces/torch_interface.py +476 -0
  288. nvmath/sparse/ust/semiring_ops.py +71 -0
  289. nvmath/sparse/ust/tensor.py +1027 -0
  290. nvmath/sparse/ust/tensor_drawer.py +552 -0
  291. nvmath/sparse/ust/tensor_emitter.py +558 -0
  292. nvmath/sparse/ust/tensor_format.py +935 -0
  293. nvmath/sparse/ust/tensor_jit.py +212 -0
  294. nvmath/sparse/ust/tensor_kernel.py +348 -0
  295. nvmath/sparse/ust/tensor_utils.py +380 -0
  296. nvmath/tensor/__init__.py +6 -0
  297. nvmath/tensor/_configuration.py +95 -0
  298. nvmath/tensor/_internal/__init__.py +3 -0
  299. nvmath/tensor/_internal/cutensor_config_ifc.py +247 -0
  300. nvmath/tensor/_internal/cutensor_utils.py +162 -0
  301. nvmath/tensor/_internal/data.py +43 -0
  302. nvmath/tensor/_internal/einsum_parser.py +444 -0
  303. nvmath/tensor/_internal/typemaps.py +96 -0
  304. nvmath/tensor/contract.py +1861 -0
  305. nvmath_python-0.9.0.dist-info/METADATA +127 -0
  306. nvmath_python-0.9.0.dist-info/RECORD +309 -0
  307. nvmath_python-0.9.0.dist-info/WHEEL +5 -0
  308. nvmath_python-0.9.0.dist-info/licenses/LICENSE +177 -0
  309. nvmath_python-0.9.0.dist-info/top_level.txt +2 -0
@@ -0,0 +1,771 @@
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ """
6
+ A collection of (internal use) helper functions.
7
+ """
8
+
9
+ import contextlib
10
+ import functools
11
+ import time
12
+ import typing
13
+ from collections.abc import Callable, Sequence
14
+
15
+ import cuda.bindings.runtime as cbr
16
+
17
+ try:
18
+ from cuda.core import Device, Event, EventOptions
19
+ except ImportError:
20
+ from cuda.core.experimental import Device, Event, EventOptions
21
+
22
+ from nvmath._internal.layout import is_contiguous_and_dense
23
+
24
+ from . import formatters, mem_limit, package_wrapper
25
+ from .package_ifc import AnyStream, Stream, StreamHolder
26
+ from .tensor_ifc import AnyTensor, Tensor, TensorHolder
27
+
28
+
29
+ def infer_object_package(obj: object) -> str:
30
+ """
31
+ Infer the package that defines this object.
32
+ """
33
+ module = obj.__class__.__module__
34
+ return module.split(".")[0]
35
+
36
+
37
+ """A generic type for Options classes"""
38
+ Options = typing.TypeVar("Options")
39
+
40
+
41
+ def check_or_create_options(
42
+ cls: type[Options], options: None | dict[str, typing.Any] | Options, options_description: str, *, keep_none=False
43
+ ) -> Options | None:
44
+ """
45
+ Create the specified options dataclass from a dictionary of options or None.
46
+ """
47
+
48
+ if options is None:
49
+ if keep_none:
50
+ return options
51
+ options = cls()
52
+ elif isinstance(options, dict):
53
+ options = cls(**options)
54
+
55
+ if not isinstance(options, cls):
56
+ raise TypeError(
57
+ f"The {options_description} must be provided as an object "
58
+ f"of type {cls.__name__} or as a dict with valid {options_description}. "
59
+ f"The provided object is '{options}'."
60
+ )
61
+
62
+ return options
63
+
64
+
65
+ def check_or_create_one_of_options(clss, options, options_description, *, cls_key="name", default_name=None):
66
+ """
67
+ Create one of the specified options dataclasses by name or from a dictionary of options.
68
+ """
69
+
70
+ assert isinstance(clss, tuple)
71
+ assert all(hasattr(cls, "name") and isinstance(cls.name, str) for cls in clss)
72
+
73
+ if options is None:
74
+ assert default_name is not None, "Internal error in execution options"
75
+ options = _create_one_of_options_from_name(clss, default_name, options_description, cls_key=cls_key)
76
+ elif isinstance(options, str):
77
+ options = _create_one_of_options_from_name(clss, options, options_description, cls_key=cls_key)
78
+ elif isinstance(options, dict):
79
+ options = _create_one_of_options_from_dict(
80
+ clss,
81
+ options,
82
+ options_description,
83
+ cls_key=cls_key,
84
+ default_name=default_name,
85
+ )
86
+ else:
87
+ if not isinstance(options, clss):
88
+ _raise_invalid_one_of_options(clss, options, options_description, cls_key=cls_key)
89
+ return options
90
+
91
+
92
+ def _create_one_of_options_from_name(clss, cls_name, options_description, *, cls_key="name"):
93
+ try:
94
+ _cls_name = cls_name.lower()
95
+ return next(cls for cls in clss if cls.name == _cls_name)()
96
+ except StopIteration:
97
+ raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key) from None
98
+
99
+
100
+ def _create_one_of_options_from_dict(
101
+ clss,
102
+ options,
103
+ options_description,
104
+ *,
105
+ cls_key="name",
106
+ default_name=None,
107
+ ):
108
+ cls_name = options.get(cls_key)
109
+ if cls_name is None:
110
+ cls_name = default_name
111
+ if not isinstance(cls_name, str):
112
+ raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key)
113
+ try:
114
+ _cls_name = cls_name.lower()
115
+ cls = next(cls for cls in clss if cls.name == _cls_name)
116
+ except StopIteration:
117
+ raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key) from None
118
+ return cls(**{key: name for key, name in options.items() if key != cls_key})
119
+
120
+
121
+ def _raise_invalid_one_of_options(clss, options, options_description, *, cls_key="name"):
122
+ accepted_names = ", ".join(f"'{cls.name}'" for cls in clss)
123
+ accepted_types = ", ".join(str(cls) for cls in clss)
124
+ raise ValueError(
125
+ f"The {options_description} must be: \n"
126
+ f"1. an object of one of the following types {accepted_types} or\n"
127
+ f"2. a string with the type name ({accepted_names}) or\n"
128
+ f"3. a dict with a '{cls_key}' key and one of the type names {accepted_names}' as a value, "
129
+ f"and optional options valid for that type. \n"
130
+ f"The provided object is '{options}'."
131
+ )
132
+
133
+
134
+ class device_ctx:
135
+ """
136
+ NOTE, using classic class-based context manager here as it has lower overhead
137
+ than the generator-based context manager.
138
+
139
+ Semantics:
140
+
141
+ 1. The device context manager makes the specified device current from the point of entry
142
+ until the point of exit.
143
+
144
+ 2. When the context manager exits, the current device is reset to what it was when the
145
+ context manager was entered.
146
+
147
+ 3. Any explicit setting of the device within the context manager (using
148
+ cupy.cuda.Device().use(), torch.cuda.set_device(), etc) will overrule the device set
149
+ by the context manager from that point onwards till the context manager exits. In
150
+ other words, the context manager provides a local device scope and the current device
151
+ can be explicitly reset for the remainder of that scope.
152
+
153
+ 4. The context manager is single-use.
154
+
155
+ Corollary: if any library function resets the device globally and this is an undesired
156
+ side-effect, such functions must be called from within the device context manager.
157
+
158
+ Device context managers can be arbitrarily nested.
159
+ """
160
+
161
+ __slots__ = ("new_device_id", "_old_device")
162
+
163
+ def __init__(self, new_device_id):
164
+ self.new_device_id = new_device_id
165
+ self._old_device = None
166
+
167
+ def __enter__(self):
168
+ if self._old_device is not None:
169
+ raise RuntimeError("Reusing a device_ctx instance is not allowed.")
170
+ self._old_device = old_device = Device()
171
+ new_device_id = self.new_device_id
172
+ if old_device.device_id != new_device_id:
173
+ device = Device(new_device_id)
174
+ device.set_current()
175
+ return device
176
+ else:
177
+ return old_device
178
+
179
+ def __exit__(self, type, value, traceback):
180
+ # We should always restore the old device at exit.
181
+ self._old_device.set_current()
182
+
183
+
184
+ def is_hashable(obj: object) -> bool:
185
+ try:
186
+ hash(obj)
187
+ except TypeError:
188
+ return False
189
+ return True
190
+
191
+
192
+ @functools.lru_cache(maxsize=128)
193
+ def cached_get_or_create_stream(
194
+ device_id: int,
195
+ stream: Stream | int,
196
+ op_package: str,
197
+ ) -> StreamHolder[Stream]:
198
+ op_package_ifc = package_wrapper.PACKAGE[op_package]
199
+ if isinstance(stream, int):
200
+ ptr = stream
201
+ external = op_package_ifc.create_external_stream(device_id, ptr)
202
+ ctx = op_package_ifc.to_stream_context(external)
203
+ obj = op_package_ifc.create_stream(external, device_id)
204
+ return StreamHolder(ctx=ctx, device_id=device_id, external=external, obj=obj, ptr=ptr, package=op_package)
205
+
206
+ stream_package = infer_object_package(stream)
207
+ if stream_package != op_package:
208
+ message = (
209
+ f"The stream object must belong to the same package as the tensor network operands. "
210
+ f"Stream package: {stream_package}, Tensor package: {op_package}"
211
+ )
212
+ raise TypeError(message)
213
+
214
+ ctx = op_package_ifc.to_stream_context(stream)
215
+ ptr = op_package_ifc.to_stream_pointer(stream)
216
+ external = stream
217
+ obj = op_package_ifc.create_stream(stream, device_id)
218
+ return StreamHolder(ctx=ctx, device_id=device_id, external=external, obj=obj, ptr=ptr, package=op_package)
219
+
220
+
221
+ def get_or_create_stream(
222
+ device_id: int,
223
+ stream: Stream | int | None,
224
+ op_package: str,
225
+ ) -> StreamHolder[Stream]:
226
+ """
227
+ Create a stream object from a stream pointer or extract the stream pointer from a stream
228
+ object, or use the current stream.
229
+
230
+ Args:
231
+ device_id: The device ID.
232
+
233
+ stream: A stream object, stream pointer, or None.
234
+
235
+ op_package: The package the tensor network operands belong to.
236
+
237
+ Returns:
238
+ StreamHolder: Hold a CuPy stream object, package stream context, stream pointer, ...
239
+ """
240
+ op_package_ifc = package_wrapper.PACKAGE[op_package]
241
+ if stream is None:
242
+ stream = op_package_ifc.get_current_stream(device_id)
243
+ return cached_get_or_create_stream(device_id, stream, op_package)
244
+
245
+
246
+ @functools.lru_cache(maxsize=128)
247
+ def get_memory_limit_from_device_id(memory_limit: int | float | str, device_id: int) -> int:
248
+ with device_ctx(device_id):
249
+ status, _, total_memory = cbr.cudaMemGetInfo()
250
+ if status != 0 or total_memory is None:
251
+ raise RuntimeError(
252
+ f"cudaMemGetInfo failed with status {status}, total_memory={total_memory}. A possible cause is \
253
+ an inconsistent version of cuda-bindings for the CTK version used with nvmath-python."
254
+ )
255
+ return _get_memory_limit(memory_limit, total_memory)
256
+
257
+
258
+ def _get_memory_limit(memory_limit: int | float | str, total_memory: int) -> int:
259
+ """
260
+ Parse user provided memory limit and return the memory limit in bytes.
261
+ """
262
+ if isinstance(memory_limit, int):
263
+ if memory_limit < 0:
264
+ raise ValueError("The specified memory limit must be greater than or equal to 0.")
265
+ return memory_limit
266
+
267
+ if isinstance(memory_limit, float):
268
+ if memory_limit < 0:
269
+ raise ValueError("The specified memory limit must be greater than or equal to 0.")
270
+ if memory_limit <= 1.0:
271
+ memory_limit *= total_memory
272
+ return int(memory_limit)
273
+
274
+ m = mem_limit.MEM_LIMIT_RE_PCT.match(memory_limit)
275
+ if m:
276
+ factor = float(m.group(1))
277
+ if factor <= 0 or factor > 100:
278
+ raise ValueError("The memory limit percentage must be in the range [0, 100].")
279
+ return int(factor * total_memory / 100.0)
280
+
281
+ m = mem_limit.MEM_LIMIT_RE_VAL.match(memory_limit)
282
+ if not m:
283
+ raise ValueError(mem_limit.MEM_LIMIT_DOC.format(kind="memory limit", value=memory_limit))
284
+
285
+ base = 1000
286
+ if m.group("binary"):
287
+ base = 1024
288
+
289
+ powers = {"": 0, "k": 1, "m": 2, "g": 3}
290
+ unit = m.group("units").lower() if m.group("units") else ""
291
+ multiplier = base ** powers[unit]
292
+
293
+ value = float(m.group("value"))
294
+ memory_limit = int(value * multiplier)
295
+
296
+ return memory_limit
297
+
298
+
299
+ def create_empty_tensor(
300
+ cls: type[TensorHolder[Tensor]],
301
+ extents: Sequence[int],
302
+ dtype: str,
303
+ device_id: int | typing.Literal["cpu"],
304
+ stream_holder: StreamHolder | None,
305
+ verify_strides: bool,
306
+ strides: Sequence[int] | None = None,
307
+ **context,
308
+ ) -> TensorHolder[Tensor]:
309
+ """
310
+ Create a wrapped tensor of the same type as (the wrapped) cls on the specified device
311
+ having the specified extents and dtype.
312
+
313
+ The tensor is created within a stream context to allow for asynchronous memory
314
+ allocators like CuPy's MemoryAsyncPool.
315
+
316
+ Note, the function assumes the `strides` are dense (possibly permuted).
317
+ Otherwise, the behaviour is framework specific and tensor creation may fail
318
+ or created tensor may be corrupted. Set `verify_strides` to True to check
319
+ the layout and drop the strides if the layout is not dense.
320
+ """
321
+ assert isinstance(device_id, int) or device_id == "cpu", f"Internal Error: device_id must be int or 'cpu'; not {device_id}"
322
+ # if device id is none the stream holder must be too
323
+ assert device_id != "cpu" or stream_holder is None
324
+ if strides is not None and verify_strides and not is_contiguous_and_dense(extents, strides):
325
+ strides = None
326
+ tensor = cls.empty(extents, dtype=dtype, device_id=device_id, strides=strides, stream_holder=stream_holder, **context)
327
+ return tensor
328
+
329
+
330
+ def get_operands_device_id(operands: Sequence[TensorHolder[AnyTensor]]):
331
+ """
332
+ Return the id (ordinal) of the device the operands are on, or None if it is on the CPU.
333
+ """
334
+ device_id = operands[0].device_id
335
+ if not all(operand.device_id == device_id for operand in operands):
336
+ devices = {operand.device_id for operand in operands}
337
+ raise ValueError(f"All operands are not on the same device. Devices = {devices}.")
338
+
339
+ return device_id
340
+
341
+
342
+ def get_operands_dtype(operands: Sequence[TensorHolder[AnyTensor]]):
343
+ """
344
+ Return the data type name of the tensors.
345
+ """
346
+ dtype = operands[0].dtype
347
+ if not all(operand.dtype == dtype for operand in operands[1:]):
348
+ dtypes = {operand.dtype for operand in operands}
349
+ raise ValueError(f"All tensors must have the same data type. Data types found = {dtypes}.")
350
+ return dtype
351
+
352
+
353
+ def get_operands_package(operands: Sequence[TensorHolder[AnyTensor]]) -> str:
354
+ """
355
+ Return the package name of the tensors.
356
+ """
357
+ package = infer_object_package(operands[0].tensor)
358
+ if not all(infer_object_package(operand.tensor) == package for operand in operands[1:]):
359
+ packages = {infer_object_package(operand.tensor) for operand in operands}
360
+ raise TypeError(f"All tensors must be from the same library package. Packages found = {packages}.")
361
+ return package
362
+
363
+
364
+ def check_operands_match(
365
+ orig_operands: Sequence[object] | object, new_operands: Sequence[object] | object, attribute: str, description: str
366
+ ):
367
+ """
368
+ Check if the specified attribute matches between the corresponding new and old operands,
369
+ and raise an exception if it doesn't.
370
+ """
371
+ if isinstance(orig_operands, Sequence):
372
+ assert isinstance(new_operands, Sequence), "Internal Error. Cannot compare operand sequence with non-sequence."
373
+ checks: list[bool] = [
374
+ getattr(o, attribute) == getattr(n, attribute) for o, n in zip(orig_operands, new_operands, strict=True)
375
+ ]
376
+
377
+ if not all(checks):
378
+ mismatch = [
379
+ f"{location}: {getattr(orig_operands[location], attribute)} => {getattr(new_operands[location], attribute)}"
380
+ for location, predicate in enumerate(checks)
381
+ if predicate is False
382
+ ]
383
+ mismatch = formatters.array2string(mismatch)
384
+ message = f"""\
385
+ The {description} of each new operand must match the {description} of the corresponding original operand.
386
+
387
+ The mismatch in {description} as a sequence of "position: original {description} => new {description}" is: \n{mismatch}"""
388
+ raise ValueError(message)
389
+ else:
390
+ assert isinstance(new_operands, object)
391
+ check: bool = getattr(orig_operands, attribute) == getattr(new_operands, attribute)
392
+ if not check:
393
+ message = f"""The {description} of the new operand must match the {description} of the original operand."""
394
+ raise ValueError(message)
395
+
396
+
397
+ def check_attribute_match(orig_attribute: object, new_attribute: object, description: str):
398
+ """
399
+ Check if the specified attribute matches between the corresponding new and old operands,
400
+ and raise an exception if it doesn't.
401
+ """
402
+ check = orig_attribute == new_attribute
403
+ if not check:
404
+ message = f"""The {description} of the new operand must match the {description} of the original operand."""
405
+ raise ValueError(message)
406
+
407
+
408
+ def get_ptr_from_memory_pointer(mem_ptr):
409
+ """
410
+ Access the value associated with one of the attributes 'device_ptr', 'device_pointer',
411
+ 'ptr'.
412
+
413
+ This is a compatibility function so that external memory pointer implementations from
414
+ numba and cupy can be used instead of MemoryPointer from nvmath.memory.
415
+ """
416
+ attributes = (
417
+ "device_ptr", # nvmath.memory
418
+ "device_pointer",
419
+ "ptr", # cupy.cuda.MemoryPointer
420
+ "pointer", # numba.cuda.MemoryPointer
421
+ )
422
+ for attr in attributes:
423
+ if hasattr(mem_ptr, attr):
424
+ return getattr(mem_ptr, attr)
425
+
426
+ message = f"Memory pointer objects should have one of the following attributes specifying the device pointer: {attributes}"
427
+ raise AttributeError(message)
428
+
429
+
430
+ class Value:
431
+ """
432
+ A simple value wrapper holding a default value.
433
+ """
434
+
435
+ def __init__(self, default, *, validator: Callable[[object], bool]):
436
+ """
437
+ Args:
438
+ default: The default value to use.
439
+ validator: A callable that validates the provided value.
440
+ """
441
+ self.validator = validator
442
+ self._data = default
443
+
444
+ @property
445
+ def data(self):
446
+ return self._data
447
+
448
+ @data.setter
449
+ def data(self, value):
450
+ self._data = self._validate(value)
451
+
452
+ def _validate(self, value):
453
+ if self.validator(value):
454
+ return value
455
+ raise ValueError(f"Internal Error: value '{value}' is not valid.")
456
+
457
+
458
+ @contextlib.contextmanager
459
+ def cuda_call_ctx(stream_holder: StreamHolder[AnyStream], blocking=True, timing=True) -> typing.Iterator[tuple[Event, Value]]:
460
+ """
461
+ A simple context manager that provides (non-)blocking behavior depending on the
462
+ `blocking` parameter for CUDA calls. The call is timed only for blocking behavior when
463
+ timing is requested.
464
+
465
+ An `end` event is recorded after the CUDA call for use in establishing stream ordering
466
+ for non-blocking calls. This event is returned together with a `Value` object that
467
+ stores the elapsed time if the call is blocking and timing is requested, or None
468
+ otherwise.
469
+
470
+ Inside the context, the current device is the device which owns the provided stream.
471
+ """
472
+ device_id = stream_holder.device_id
473
+ stream = stream_holder.obj
474
+
475
+ with device_ctx(device_id):
476
+ if timing:
477
+ start = stream.record(options=EventOptions(enable_timing=blocking))
478
+
479
+ end = stream.device.create_event(options=EventOptions(enable_timing=(timing and blocking)))
480
+
481
+ time = Value(None, validator=lambda v: True)
482
+ yield end, time
483
+
484
+ stream.record(end)
485
+
486
+ if not blocking:
487
+ return
488
+
489
+ end.sync()
490
+
491
+ if timing:
492
+ time.data = end - start
493
+
494
+
495
+ @contextlib.contextmanager
496
+ def host_call_ctx(timing=False):
497
+ elapsed = Value(None, validator=lambda v: True)
498
+
499
+ if timing:
500
+ start_time = time.perf_counter_ns()
501
+
502
+ yield elapsed
503
+
504
+ if timing:
505
+ elapsed.data = (time.perf_counter_ns() - start_time) * 1e-6
506
+
507
+
508
+ # Decorator definitions
509
+
510
+
511
+ def atomic(
512
+ handler: Callable[[typing.Any, Exception | None], bool] | Callable[[Exception | None], bool], method: bool = False
513
+ ) -> Callable:
514
+ """
515
+ A decorator that provides "succeed or roll-back" semantics. A typical use for this is to
516
+ release partial resources if an exception occurs.
517
+
518
+ Args:
519
+ handler: A function to call when an exception occurs. The handler takes a single
520
+ argument, which is the exception object, and returns a boolean stating whether
521
+ the same exception should be reraised. We assume that this function does not
522
+ raise an exception.
523
+
524
+ method: Specify if the wrapped function as well as the exception handler are methods
525
+ bound to the same object (method = True) or they are free functions (method =
526
+ False).
527
+
528
+ Returns:
529
+ Callable: A decorator that creates the wrapping.
530
+ """
531
+
532
+ def outer(wrapped_function):
533
+ """
534
+ A decorator that actually wraps the function for exception handling.
535
+ """
536
+
537
+ @functools.wraps(wrapped_function)
538
+ def inner(*args, **kwargs):
539
+ """
540
+ Call the wrapped function and return the result. If an exception occurs, then
541
+ call the exception handler and reraise the exception.
542
+ """
543
+ try:
544
+ result = wrapped_function(*args, **kwargs)
545
+ except BaseException as e:
546
+ if method:
547
+ flag = handler(args[0], e)
548
+ else:
549
+ flag = handler(e)
550
+
551
+ if flag:
552
+ raise e
553
+
554
+ return result
555
+
556
+ return inner
557
+
558
+ return outer
559
+
560
+
561
+ def precondition(checker: Callable[..., None], what: str = "") -> Callable:
562
+ """
563
+ A decorator that adds checks to ensure any preconditions are met.
564
+
565
+ Args:
566
+ checker: The function to call to check whether the preconditions are met. It has the
567
+ same signature as the wrapped function with the addition of the keyword argument
568
+ `what`.
569
+ what: A string that is passed in to `checker` to provide context information.
570
+
571
+ Returns:
572
+ Callable: A decorator that creates the wrapping.
573
+ """
574
+
575
+ def outer(wrapped_function):
576
+ """
577
+ A decorator that actually wraps the function for checking preconditions.
578
+ """
579
+
580
+ @functools.wraps(wrapped_function)
581
+ def inner(*args, **kwargs):
582
+ """
583
+ Check preconditions and if they are met, call the wrapped function.
584
+ """
585
+ checker(*args, **kwargs, what=what)
586
+ result = wrapped_function(*args, **kwargs)
587
+
588
+ return result
589
+
590
+ return inner
591
+
592
+ return outer
593
+
594
+
595
+ COMMON_SHARED_DOC_MAP = {
596
+ "operand": """\
597
+ A tensor (ndarray-like object). The currently supported types are :class:`numpy.ndarray`,
598
+ :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
599
+ #
600
+ "stream": """\
601
+ Provide the CUDA stream to use for executing the operation. Acceptable inputs include
602
+ ``cudaStream_t`` (as Python :class:`int`), :class:`cupy.cuda.Stream`, and
603
+ :class:`torch.cuda.Stream`. If a stream is not provided, the current stream from the operand
604
+ package will be used. See :ref:`stream-semantics-guide` for more details on stream
605
+ handling.""".replace("\n", " "),
606
+ #
607
+ "release_workspace": """\
608
+ A value of `True` specifies that the stateful object should release workspace memory back to
609
+ the package memory pool on function return, while a value of `False` specifies that the
610
+ object should retain the memory. This option may be set to `True` if the application
611
+ performs other operations that consume a lot of memory between successive calls to the (same
612
+ or different) :meth:`execute` API, but incurs a small overhead due to obtaining and
613
+ releasing workspace memory from and to the package memory pool on every call. The default is
614
+ `False`.""".replace("\n", " "),
615
+ }
616
+
617
+
618
+ def _release_operand_docstring(plural, version_added=None):
619
+ operands = "operands" if plural else "operand"
620
+ their = "their" if plural else "its"
621
+ counts = "counts" if plural else "count"
622
+ reside = "reside" if plural else "resides"
623
+ reset_method = "reset_operands" if plural else "reset_operand"
624
+ reset_method_unchecked = "reset_operands_unchecked" if plural else "reset_operand_unchecked"
625
+ new_operands = "new operands" if plural else "a new operand"
626
+ are = "are" if plural else "is"
627
+ version_line = f"\n\n .. versionadded:: {version_added}" if version_added else ""
628
+ return f"""\
629
+ .. experimental:: method
630
+ {version_line}
631
+
632
+ This method does two things:
633
+
634
+ - Releases internal references to the user-provided {operands},
635
+ so that this instance no longer contributes to {their} reference {counts}.
636
+
637
+ - Frees any internal copies (mirrors) that were created when the
638
+ user-provided {operands} {reside} in a different memory space than
639
+ the execution (i.e., copies made during construction or
640
+ :meth:`{reset_method}` / :meth:`{reset_method_unchecked}`
641
+ if present).
642
+
643
+ This functionality can be useful in memory-constrained scenarios, e.g. where
644
+ multiple stateful objects need to coexist. Leveraging this functionality,
645
+ the caller can reduce memory usage while retaining the planned state.
646
+
647
+ Args:
648
+ None
649
+
650
+ Returns:
651
+ None
652
+
653
+ Semantics:
654
+ - Preserves the planned state of the stateful object.
655
+
656
+ - After calling this method, :meth:`{reset_method}` (or
657
+ :meth:`{reset_method_unchecked}` if present) must be called to
658
+ supply {new_operands} before the next :meth:`execute` call.
659
+ Failure to do so will result in a runtime error.
660
+ Device-side copies will be re-allocated as needed.
661
+
662
+ - For cross-space scenarios (e.g. CPU {operands} with GPU execution,
663
+ or GPU {operands} with CPU execution): execution is guaranteed to be
664
+ always blocking, so :meth:`execute` does not return until all
665
+ computation is complete. It is therefore always safe to call this
666
+ method after calling :meth:`execute` without additional synchronization.
667
+
668
+ - When the {operands} {are} in the same memory space as the execution
669
+ (e.g. GPU {operands} with GPU execution): in such case, this method
670
+ drops this instance's internal reference to the user-provided {operands}.
671
+ If the reference count of the {operands} reaches zero, {their} memory
672
+ may be freed, so particular attention should be paid. The caller is
673
+ responsible to ensure that if such deallocation happens, it is
674
+ ordered after pending computation (e.g. by retaining a reference until
675
+ the computation is complete, or by synchronizing the stream).
676
+ Failure to do so is analogous to use-after-free.
677
+
678
+ See :ref:`nvmath overview`, :ref:`stateful_apis_guide` for operand
679
+ lifecycle and usage patterns, and :ref:`stream-semantics-guide`
680
+ for stream ordering rules.
681
+ """
682
+
683
+
684
+ COMMON_SHARED_DOC_MAP["release_operands"] = _release_operand_docstring(True, version_added="0.9.0")
685
+ COMMON_SHARED_DOC_MAP["release_operand"] = _release_operand_docstring(False, version_added="0.9.0")
686
+
687
+
688
+ def _reset_operand_unchecked_docstring(plural, version_added=None, validation_examples="package match, data type match"):
689
+ operands = "operands" if plural else "operand"
690
+ reset_method = "reset_operands" if plural else "reset_operand"
691
+ release_method = "release_operands" if plural else "release_operand"
692
+ version_line = f"\n\n .. versionadded:: {version_added}" if version_added else ""
693
+ return f"""\
694
+ .. experimental:: method
695
+ {version_line}
696
+
697
+ This method is a performance-optimized alternative to :meth:`{reset_method}`
698
+ that eliminates validation and logging overhead, making it ideal for
699
+ performance-critical loops where {operands}
700
+ compatibility is guaranteed by the caller.
701
+
702
+ This method accepts the same parameters as :meth:`{reset_method}`.
703
+
704
+ Semantics:
705
+ The semantics are the same as in :meth:`{reset_method}`,
706
+ except that this method does not perform any validation
707
+ (e.g. {validation_examples}, etc.) or logging.
708
+
709
+ When to Use:
710
+ - Performance-critical loops with repeated executions on different {operands}
711
+
712
+ - After verifying correctness with :meth:`{reset_method}` during development
713
+
714
+ - When {operands} compatibility is guaranteed by construction or invariant
715
+
716
+ .. seealso::
717
+ :meth:`{reset_method}`, :meth:`{release_method}`
718
+ """
719
+
720
+
721
+ COMMON_SHARED_DOC_MAP["reset_operands_unchecked"] = _reset_operand_unchecked_docstring(True, version_added="0.9.0")
722
+ COMMON_SHARED_DOC_MAP["reset_operand_unchecked"] = _reset_operand_unchecked_docstring(False, version_added="0.9.0")
723
+
724
+
725
+ class DefaultDocstring(dict):
726
+ def __missing__(self, key):
727
+ return "{" + key + "}"
728
+
729
+
730
+ def docstring_decorator(doc_map, skip_missing=False):
731
+ assert isinstance(doc_map, dict)
732
+
733
+ def _format_doc(doc):
734
+ if doc is not None:
735
+ if skip_missing:
736
+ # Using a default class to handle missing keys in doc_map
737
+ doc = doc.format_map(DefaultDocstring(**doc_map))
738
+ else:
739
+ doc = doc.format(**doc_map)
740
+ return doc
741
+
742
+ def decorator(func_or_class):
743
+ if isinstance(func_or_class, type): # class decorator
744
+ # update the docstring of all public methods with docstrings
745
+ static_methods = [] # staticmethods appear to require special handling
746
+ for name, method in vars(func_or_class).items():
747
+ if isinstance(method, (staticmethod, classmethod)):
748
+ static_methods.append(name)
749
+ continue
750
+ if isinstance(method, property):
751
+ if method.fget and method.fget.__doc__:
752
+ # Create a new property with updated docstring
753
+ setattr(
754
+ func_or_class,
755
+ name,
756
+ property(method.fget, method.fset, method.fdel, _format_doc(method.fget.__doc__)),
757
+ )
758
+ continue
759
+ if callable(method) and (not name.startswith("_")) and method.__doc__:
760
+ method.__doc__ = _format_doc(method.__doc__)
761
+ # update the docstring of the constructor
762
+ func_or_class.__doc__ = _format_doc(func_or_class.__doc__)
763
+ for name in static_methods:
764
+ method = func_or_class.__dict__[name]
765
+ method.__func__.__doc__ = _format_doc(method.__func__.__doc__)
766
+ return func_or_class
767
+ else: # function decorator
768
+ func_or_class.__doc__ = _format_doc(func_or_class.__doc__)
769
+ return func_or_class
770
+
771
+ return decorator