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,935 @@
1
+ # Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ """
6
+ The tensor format DSL of the Universal Sparse Tensor maps tensor dimensions
7
+ (logical tensor indices) to storage levels (physical memory indices) using
8
+ an invertible function that defines how each level should be stored.
9
+ """
10
+
11
+ __all__ = [
12
+ "is_ordered",
13
+ "is_unique",
14
+ "Dimension",
15
+ "LevelFormat",
16
+ "LevelProperty",
17
+ "LevelExpr",
18
+ "TensorFormat",
19
+ "NamedFormats",
20
+ ]
21
+
22
+ from collections.abc import Sequence
23
+ from enum import IntEnum
24
+
25
+
26
+ def is_ordered(prop):
27
+ """Determines if given property is ordered."""
28
+ if isinstance(prop, Sequence):
29
+ for p in prop:
30
+ if not is_ordered(p):
31
+ return False
32
+ elif isinstance(prop, LevelProperty) and prop == LevelProperty.UNORDERED:
33
+ return False
34
+ return True # ordered unless seen otherwise
35
+
36
+
37
+ def is_unique(prop):
38
+ """Determines if given property is unique."""
39
+ if isinstance(prop, Sequence):
40
+ for p in prop:
41
+ if not is_unique(p):
42
+ return False
43
+ elif isinstance(prop, LevelProperty) and prop == LevelProperty.NONUNIQUE:
44
+ return False
45
+ return True # unique unless seen otherwise
46
+
47
+
48
+ class LevelFormat(IntEnum):
49
+ DENSE = 0
50
+ BATCH = 1
51
+ COMPRESSED = 2
52
+ SINGLETON = 3
53
+ RANGE = 4
54
+ DELTA = 5
55
+
56
+ def __repr__(self):
57
+ """Returns representation without index."""
58
+ return f"<{self.__class__.__name__}.{self.name}>"
59
+
60
+
61
+ class LevelProperty(IntEnum):
62
+ UNIQUE = 0
63
+ NONUNIQUE = 1
64
+ ORDERED = 2
65
+ UNORDERED = 3
66
+
67
+ def __repr__(self):
68
+ """Returns representation without index."""
69
+ return f"<{self.__class__.__name__}.{self.name}>"
70
+
71
+
72
+ # The inverse map is used to check for certain incompatible property sequences.
73
+ LevelPropertyInverse = {
74
+ LevelProperty.UNIQUE: LevelProperty.NONUNIQUE,
75
+ LevelProperty.NONUNIQUE: LevelProperty.UNIQUE,
76
+ LevelProperty.ORDERED: LevelProperty.UNORDERED,
77
+ LevelProperty.UNORDERED: LevelProperty.ORDERED,
78
+ }
79
+
80
+
81
+ def _find_range(levels, expr1, expr2):
82
+ for rl, (k, v) in enumerate(levels.items()):
83
+ if v == LevelFormat.RANGE and k in (expr1, expr2):
84
+ return rl, k
85
+ raise ValueError(f"Cannot find {expr1} or {expr2} in levels")
86
+
87
+
88
+ class Operator:
89
+ """
90
+ An abstract binary level expression operator type.
91
+ """
92
+
93
+ # TODO: use protocol.
94
+ pass
95
+
96
+
97
+ class Add(Operator):
98
+ """
99
+ Add two level expressions.
100
+ """
101
+
102
+ def __call__(self, expr1, expr2, as_size):
103
+ if as_size:
104
+ return expr1 + expr2 - 1
105
+ return expr1 + expr2
106
+
107
+ def invert(self, expr1, expr2, dimensions, levels, dim_indices, lvl_indices, idx):
108
+ rl, re = _find_range(levels, expr1, expr2)
109
+ if re == expr2:
110
+ dim_indices[dimensions.index(expr1)] = lvl_indices[idx] - lvl_indices[rl]
111
+ else:
112
+ assert re == expr1
113
+ dim_indices[dimensions.index(expr2)] = lvl_indices[idx] - lvl_indices[rl]
114
+
115
+ def __repr__(self):
116
+ return "+"
117
+
118
+
119
+ class Subtract(Operator):
120
+ """
121
+ Subtract a level expression from another.
122
+ """
123
+
124
+ def __call__(self, expr1, expr2, as_size):
125
+ if as_size:
126
+ return expr1 + expr2 - 1
127
+ return expr1 - expr2
128
+
129
+ def invert(self, expr1, expr2, dimensions, levels, dim_indices, lvl_indices, idx):
130
+ rl, re = _find_range(levels, expr1, expr2)
131
+ if re == expr2:
132
+ dim_indices[dimensions.index(expr1)] = lvl_indices[rl] + lvl_indices[idx]
133
+ else:
134
+ assert re == expr1
135
+ dim_indices[dimensions.index(expr2)] = lvl_indices[rl] - lvl_indices[idx]
136
+
137
+ def __repr__(self):
138
+ return "-"
139
+
140
+
141
+ class Divide(Operator):
142
+ """
143
+ Divide a level expression by an integer.
144
+ """
145
+
146
+ def __call__(self, expr1, expr2, as_size):
147
+ if as_size:
148
+ assert expr2 > 0, f"Internal error: {expr2}"
149
+ if expr1 % expr2 != 0:
150
+ raise NotImplementedError(f"block size {expr2} does not evenly divide size {expr1}")
151
+ return expr1 // expr2
152
+
153
+ def invert(self, expr1, expr2, dimensions, levels, dim_indices, lvl_indices, idx):
154
+ i = dimensions.index(expr1)
155
+ dim_indices[i] = lvl_indices[idx] * expr2
156
+
157
+ def __repr__(self):
158
+ return "//"
159
+
160
+
161
+ class Modulo(Operator):
162
+ """
163
+ The integer modulo of a level expression.
164
+ """
165
+
166
+ def __call__(self, expr1, expr2, as_size):
167
+ if as_size:
168
+ assert expr2 > 0, f"Internal error: {expr2}"
169
+ return expr2
170
+ return expr1 % expr2
171
+
172
+ def invert(self, expr1, expr2, dimensions, levels, dim_indices, lvl_indices, idx):
173
+ i = dimensions.index(expr1)
174
+ dim_indices[i] += lvl_indices[idx] # update (seen after div)
175
+
176
+ def __repr__(self):
177
+ return "%"
178
+
179
+
180
+ class LevelExpr:
181
+ """
182
+ A (binary) level expression is a triple: dimension object or level expression,
183
+ operator, dimension or level expression or int.
184
+
185
+ Args:
186
+ expression1: a Dimension or LevelExpr object.
187
+ operator: operator derived from Operator (Add, Subtract, Divide, Modulo).
188
+ expression2: a Dimension or LevelExpr object, or an integer.
189
+ """
190
+
191
+ __slots__ = ("expression1", "operator", "expression2")
192
+
193
+ def __init__(self, *, expression1, operator, expression2):
194
+ lhs_level_expr_types = Dimension, LevelExpr
195
+ rhs_level_expr_types = lhs_level_expr_types + (int,)
196
+ if not isinstance(expression1, lhs_level_expr_types):
197
+ raise TypeError(f"LHS expression {expression1} must have one of these types: {lhs_level_expr_types}.")
198
+ if not isinstance(operator, Operator):
199
+ raise TypeError(f"Operator {operator} must be derived from Operator.")
200
+ if not isinstance(expression2, rhs_level_expr_types):
201
+ raise TypeError(f"RHS expression {expression2} must have one of these types: {rhs_level_expr_types}.")
202
+ self.expression1 = expression1
203
+ self.operator = operator
204
+ self.expression2 = expression2
205
+
206
+ def __repr__(self):
207
+ return f"({self.expression1} {self.operator} {self.expression2})"
208
+
209
+ def __add__(self, expression):
210
+ return LevelExpr(expression1=self, operator=Add(), expression2=expression)
211
+
212
+ def __sub__(self, expression):
213
+ return LevelExpr(expression1=self, operator=Subtract(), expression2=expression)
214
+
215
+ def __floordiv__(self, divisor):
216
+ return LevelExpr(expression1=self, operator=Divide(), expression2=divisor)
217
+
218
+ def __mod__(self, modulus):
219
+ return LevelExpr(expression1=self, operator=Modulo(), expression2=modulus)
220
+
221
+ def evaluate(self, dimensions, dim_indices, as_size):
222
+ e1 = self.expression1.evaluate(dimensions, dim_indices, as_size)
223
+ e2 = (
224
+ self.expression2.evaluate(dimensions, dim_indices, as_size)
225
+ if not isinstance(self.expression2, int)
226
+ else self.expression2
227
+ )
228
+ return self.operator(e1, e2, as_size)
229
+
230
+ def invert(self, dimensions, levels, dim_indices, lvl_indices, idx):
231
+ self.operator.invert(
232
+ self.expression1,
233
+ self.expression2,
234
+ dimensions,
235
+ levels,
236
+ dim_indices,
237
+ lvl_indices,
238
+ idx,
239
+ )
240
+
241
+
242
+ class Dimension:
243
+ """
244
+ A dimension object encapsulates a dimension name.
245
+
246
+ Args:
247
+ dimension_name: A name for the dimension.
248
+ """
249
+
250
+ __slots__ = ("dimension_name",)
251
+
252
+ def __init__(self, *, dimension_name):
253
+ if not isinstance(dimension_name, str):
254
+ raise TypeError(f"The dimension name {dimension_name} must be a string.")
255
+
256
+ self.dimension_name = dimension_name
257
+
258
+ def __repr__(self):
259
+ return f"{self.dimension_name}"
260
+
261
+ def __add__(self, expression):
262
+ return LevelExpr(expression1=self, operator=Add(), expression2=expression)
263
+
264
+ def __sub__(self, expression):
265
+ return LevelExpr(expression1=self, operator=Subtract(), expression2=expression)
266
+
267
+ def __floordiv__(self, divisor):
268
+ return LevelExpr(expression1=self, operator=Divide(), expression2=divisor)
269
+
270
+ def __mod__(self, modulus):
271
+ return LevelExpr(expression1=self, operator=Modulo(), expression2=modulus)
272
+
273
+ def evaluate(self, dimensions, dim_indices, as_size):
274
+ return dim_indices[dimensions.index(self)]
275
+
276
+ def invert(self, dimensions, levels, dim_indices, lvl_indices, idx):
277
+ dim_indices[dimensions.index(self)] = lvl_indices[idx]
278
+
279
+
280
+ class TensorFormat:
281
+ """
282
+ A universal sparse tensor format maps dimension specifications
283
+ (dimensions for short) to level specifications (levels for short).
284
+
285
+ Args:
286
+ dimensions: a sequence of Dimension objects.
287
+ levels: an ordered dictionary from a LevelExpr or Dimension object
288
+ to a LevelFormat, or a (LevelFormat, LevelProperty) pair.
289
+ name: a name for the format as a string. If the format corresponds to an
290
+ existing format like CSR, COO, etc. use the canonical name. A name
291
+ will be generated based on the level specification if none is provided.
292
+ """
293
+
294
+ __slots__ = ("_dimensions", "_levels", "_name", "_is_identity", "_is_ordered", "_is_unique")
295
+
296
+ def __init__(self, dimensions, levels, *, name=None):
297
+ self._dimensions = dimensions
298
+ self._levels = levels
299
+
300
+ # TODO: create a better default name.
301
+ if name is None:
302
+ name = f"{levels}".replace("\n", " ")
303
+ if not isinstance(name, str):
304
+ raise TypeError(f"The name {name} must be a string.")
305
+ self._name = name
306
+
307
+ # Perform semantic checks on syntactic structure.
308
+ # Throws a TypeError on failure.
309
+ self._semantically_validate()
310
+
311
+ # Cached properties.
312
+ self._is_identity = (
313
+ all(isinstance(lvl, Dimension) and self.dimensions.index(lvl) == idx for idx, lvl in enumerate(levels))
314
+ if self.num_dimensions == self.num_levels
315
+ else False
316
+ )
317
+ self._is_ordered = all(is_ordered(lvl[1]) if isinstance(lvl, tuple) else True for lvl in levels.values())
318
+ self._is_unique = any(is_unique(lvl[1]) if isinstance(lvl, tuple) else True for lvl in levels.values())
319
+
320
+ @property
321
+ def name(self):
322
+ """
323
+ Get the name of the sparse tensor format.
324
+ """
325
+ return self._name
326
+
327
+ @property
328
+ def num_dimensions(self):
329
+ """
330
+ Get the number of dimensions (rank of tensor).
331
+ """
332
+ return len(self._dimensions)
333
+
334
+ @property
335
+ def num_levels(self):
336
+ """
337
+ Get the number of levels (rank of storage).
338
+ """
339
+ return len(self._levels)
340
+
341
+ @property
342
+ def dimensions(self):
343
+ """
344
+ Get the dimension specifications for this format.
345
+ """
346
+ return self._dimensions
347
+
348
+ @property
349
+ def levels(self):
350
+ """
351
+ Get the level specifications for this format.
352
+ """
353
+ return self._levels
354
+
355
+ @property
356
+ def is_identity(self):
357
+ """
358
+ Determine whether the format uses identity mapping between dimensions and levels.
359
+ """
360
+ return self._is_identity
361
+
362
+ @property
363
+ def is_ordered(self):
364
+ """
365
+ Determine whether the format is ordered.
366
+ """
367
+ return self._is_ordered
368
+
369
+ @property
370
+ def is_unique(self):
371
+ """
372
+ Determine whether the format is unique.
373
+ """
374
+ # TODO: imprecise, but works for expected cases like
375
+ # e.g. COO3 (compressed(!U), singleton(!U), singleton(U)) is still unique
376
+ # but COO3 (compressed(!U), singleton(!U), singleton(!U)) is nonunique
377
+ return self._is_unique
378
+
379
+ def _semantically_validate(self):
380
+ """
381
+ Semantically validates the sparse tensor format. The grammar allows for specifying
382
+ a very wide range of sparse tensor formats. However, not every syntactically
383
+ valid input is also semanticaly valid. This method throws a TypeError for
384
+ tensor formats where the following constraints are not satisfied :
385
+
386
+ (1) the input syntax has an actual implementation
387
+ (2) the dim-to-level mapping is easily invertible
388
+
389
+ Note that currently the checks are **stricter** than really required.
390
+ However, developers that relax the checks below are also responsible for
391
+ ensuring that constraints (1) and (2) are still satisfied.
392
+ """
393
+ # Reject incorrect dimension specifications types.
394
+ if not (isinstance(self.dimensions, Sequence) and all(isinstance(d, Dimension) for d in self.dimensions)):
395
+ raise TypeError(f"Dimension specifications {self.dimensions} must be a sequence of Dimension objects.")
396
+
397
+ # Reject dimension specifications that are not unique names (e.g. (i,i)).
398
+ all_dims = set(self.dimensions)
399
+ if len(all_dims) != len(self.dimensions):
400
+ raise TypeError(f"Dimension specifications {self.dimensions} has repeated identifiers.")
401
+
402
+ # Reject incorrect level specification types.
403
+ if not (isinstance(self.levels, dict)):
404
+ raise TypeError(f"Level specifications {self.levels} must be an ordered dictionary.")
405
+
406
+ # Reject incorrect level specification structure.
407
+ used_dims = set()
408
+ range_dims = {}
409
+ block_dims = {}
410
+ batch = 0
411
+ for idx, (k, v) in enumerate(self.levels.items()):
412
+ # Check the level expression.
413
+ if isinstance(k, Dimension):
414
+ if k not in self.dimensions:
415
+ raise TypeError(f"Dimension {k} does not appear in {self.dimensions}.")
416
+ used_dims.add(k)
417
+ elif isinstance(k, LevelExpr):
418
+ # Check for unnested level expressions and impose some constraints.
419
+ # We expect (Add, Subtract, Divide, Modulo) only.
420
+ if isinstance(k.operator, (Add, Subtract)):
421
+ # Only allow for direct expressions like i - j.
422
+ if k.expression1 not in self.dimensions:
423
+ raise TypeError(f"LHS {k.expression1} in {k} does not appear in {self.dimensions}.")
424
+ if k.expression2 not in self.dimensions:
425
+ raise TypeError(f"RHS {k.expression2} in {k} does not appear in {self.dimensions}.")
426
+ used_dims.add(k.expression1)
427
+ used_dims.add(k.expression2)
428
+ # Find matching range.
429
+ if k.expression1 in range_dims or k.expression2 in range_dims:
430
+ raise TypeError(f"Operation {k} reuses dimension of a prior range computation.")
431
+ range_dims[k.expression1] = k.expression2
432
+ range_dims[k.expression2] = k.expression1
433
+ elif isinstance(k.operator, (Divide, Modulo)):
434
+ # Only allow for direct expressions like i % 3.
435
+ if k.expression1 not in self.dimensions:
436
+ raise TypeError(f"LHS {k.expression1} in {k} does not appear in {self.dimensions}.")
437
+ if not isinstance(k.expression2, int):
438
+ raise TypeError(f"RHS {k.expression2} in {k} must be an integer.")
439
+ elif k.expression2 <= 0:
440
+ raise TypeError(f"RHS {k.expression2} in {k} must be strictly positive integer.")
441
+ used_dims.add(k.expression1)
442
+ # Find matching blocking.
443
+ if isinstance(k.operator, Divide):
444
+ if k.expression1 in block_dims:
445
+ raise TypeError(f"Division {k} reuses dimension of a prior division.")
446
+ block_dims[k.expression1] = k.expression2
447
+ elif isinstance(k.operator, Modulo):
448
+ if block_dims.get(k.expression1, 0) != k.expression2:
449
+ raise TypeError(f"Modulo {k} does not match any prior division.")
450
+ block_dims[k.expression1] = 0 # avoid further usage
451
+ else:
452
+ raise TypeError(f"Unexpected operator in level expression {k}.")
453
+ else:
454
+ raise TypeError(f"Level expression {k} must be a Dimension or LevelExpr object.")
455
+
456
+ # Inspect the level format for explicit or implicit properties.
457
+ if isinstance(v, tuple):
458
+ fmt, prop = v
459
+ else:
460
+ fmt, prop = v, LevelProperty.ORDERED
461
+
462
+ # Check level format.
463
+ if not isinstance(fmt, LevelFormat):
464
+ raise TypeError(f"Invalid level format {fmt}.")
465
+ elif fmt == LevelFormat.RANGE:
466
+ if not isinstance(k, Dimension):
467
+ raise TypeError(f"Range uses compound level expression {k}.")
468
+ other = range_dims.get(k, 0)
469
+ if other == 0:
470
+ raise TypeError(f"Range uses dimension {k} that is not uniquely defined.")
471
+ # Consume range tuple.
472
+ range_dims[k] = 0
473
+ range_dims[other] = 0
474
+
475
+ # Verify batching.
476
+ if fmt == LevelFormat.BATCH:
477
+ if k != self.dimensions[idx]:
478
+ raise TypeError(f"Batch uses non-identity {k} for {self.dimensions[idx]}.")
479
+ if batch == -1:
480
+ raise TypeError(f"Batch is used in inner level {idx}.")
481
+ batch += 1
482
+ elif batch > 0 and fmt != LevelFormat.DENSE:
483
+ raise TypeError("Batch levels must end in a dense level.")
484
+ else:
485
+ batch = -1
486
+
487
+ # Check level properties.
488
+ if isinstance(prop, Sequence):
489
+ invalid = []
490
+ for p_ in prop:
491
+ if not isinstance(p_, LevelProperty):
492
+ raise TypeError(f"Invalid level property {p_}.")
493
+ if LevelPropertyInverse[p_] in prop:
494
+ invalid += (p_,)
495
+ if invalid:
496
+ raise TypeError(f"Invalid level property combination: {invalid}.")
497
+ elif not isinstance(prop, LevelProperty) and not isinstance(prop, int):
498
+ raise TypeError(f"Invalid level property {prop}.")
499
+
500
+ # Check that every dimension is used eventually.
501
+ if len(all_dims) != len(used_dims):
502
+ raise TypeError(f"The following dimensions are not used: {all_dims - used_dims}.")
503
+
504
+ # Check that every div is matched by mod.
505
+ unmatched = [key for key, value in block_dims.items() if value != 0]
506
+ if len(unmatched) != 0:
507
+ raise TypeError(f"Some division dimensions are not matched by modulo: {unmatched}.")
508
+
509
+ # Check that every add/sub is matched by range.
510
+ unmatched = [key for key, value in range_dims.items() if value != 0]
511
+ if len(unmatched) != 0:
512
+ raise TypeError(f"Some add/sub dimensions are not matched by range: {unmatched}.")
513
+
514
+ # Check that batching ends in dense.
515
+ if batch > 0:
516
+ raise TypeError("Batch levels are not properly closed by a dense level.")
517
+
518
+ def dim2lvl(self, dim_indices, as_size=False):
519
+ """
520
+ Maps the dimension index list to the level index list.
521
+
522
+ Examples:
523
+ CSR: [0,4] -> [0,4]
524
+ CSC: [0,4] -> [4,0]
525
+ BSRRight(2,2): [1,2] -> [0,1,0]
526
+ """
527
+ return [ls.evaluate(self.dimensions, dim_indices, as_size) for ls in self.levels]
528
+
529
+ def lvl2dim(self, lvl_indices):
530
+ """
531
+ Maps the level index list to the dimension index list (the inverse of dim2lvl).
532
+ We can safely assume that the mapping is invertible given that the tensor format
533
+ passed validation.
534
+
535
+ Examples:
536
+ CSR: [0,4] -> [0,4]
537
+ CSC: [4,0] -> [0,4]
538
+ BSRRight(2,2): [0,1,0] -> [1,2]
539
+ """
540
+ dim_indices = [0] * self.num_dimensions # pre-populate
541
+ for idx, ls in enumerate(self.levels):
542
+ ls.invert(self.dimensions, self.levels, dim_indices, lvl_indices, idx)
543
+ return dim_indices
544
+
545
+ def __repr__(self):
546
+ return (
547
+ f"{self.dimensions} -> ("
548
+ + ", ".join([f"{k}: {v if isinstance(v, Sequence) else repr(v)}" for k, v in self.levels.items()])
549
+ + ")"
550
+ )
551
+
552
+
553
+ class NamedFormats:
554
+ """
555
+ A number of pre-defined common tensor formats (direct or using builders).
556
+ Clients can always construct these or other tensor formats directly.
557
+ """
558
+
559
+ i, j, k, b = (
560
+ Dimension(dimension_name="i"),
561
+ Dimension(dimension_name="j"),
562
+ Dimension(dimension_name="k"),
563
+ Dimension(dimension_name="batch"),
564
+ )
565
+
566
+ # Scalar format.
567
+ Scalar = TensorFormat([], {}, name="Scalar")
568
+
569
+ # Vector formats.
570
+ DenseVector = TensorFormat([i], {i: LevelFormat.DENSE}, name="DenseVector")
571
+ SparseVector = TensorFormat([i], {i: LevelFormat.COMPRESSED}, name="SparseVector")
572
+
573
+ # Dense matrix formats.
574
+ DenseMatrixRight = TensorFormat(
575
+ [i, j],
576
+ {i: LevelFormat.DENSE, j: LevelFormat.DENSE},
577
+ name="DenseMatrixRight",
578
+ )
579
+ DenseMatrixLeft = TensorFormat(
580
+ [i, j],
581
+ {j: LevelFormat.DENSE, i: LevelFormat.DENSE},
582
+ name="DenseMatrixLeft",
583
+ )
584
+
585
+ # Sparse matrix formats.
586
+ COO = TensorFormat(
587
+ [i, j],
588
+ {
589
+ i: (LevelFormat.COMPRESSED, LevelProperty.NONUNIQUE),
590
+ j: LevelFormat.SINGLETON,
591
+ },
592
+ name="COO",
593
+ )
594
+ CSR = TensorFormat([i, j], {i: LevelFormat.DENSE, j: LevelFormat.COMPRESSED}, name="CSR")
595
+ CSC = TensorFormat([i, j], {j: LevelFormat.DENSE, i: LevelFormat.COMPRESSED}, name="CSC")
596
+ DCSR = TensorFormat(
597
+ [i, j],
598
+ {i: LevelFormat.COMPRESSED, j: LevelFormat.COMPRESSED},
599
+ name="DCSR",
600
+ )
601
+ DCSC = TensorFormat(
602
+ [i, j],
603
+ {j: LevelFormat.COMPRESSED, i: LevelFormat.COMPRESSED},
604
+ name="DCSC",
605
+ )
606
+ CROW = TensorFormat([i, j], {i: LevelFormat.COMPRESSED, j: LevelFormat.DENSE}, name="CROW")
607
+ CCOL = TensorFormat([i, j], {j: LevelFormat.COMPRESSED, i: LevelFormat.DENSE}, name="CCOL")
608
+
609
+ # Diagonal matrix formats.
610
+ DIAI = TensorFormat(
611
+ [i, j],
612
+ {j - i: LevelFormat.COMPRESSED, i: LevelFormat.RANGE},
613
+ name="DIAI",
614
+ )
615
+ DIAJ = TensorFormat(
616
+ [i, j],
617
+ {j - i: LevelFormat.COMPRESSED, j: LevelFormat.RANGE},
618
+ name="DIAJ",
619
+ )
620
+ SkewDIAI = TensorFormat(
621
+ [i, j],
622
+ {i + j: LevelFormat.COMPRESSED, i: LevelFormat.RANGE},
623
+ name="SkewDIAI",
624
+ )
625
+ SkewDIAJ = TensorFormat(
626
+ [i, j],
627
+ {i + j: LevelFormat.COMPRESSED, j: LevelFormat.RANGE},
628
+ name="SkewDIAJ",
629
+ )
630
+
631
+ # Batched sparse matrices.
632
+ BatchedCSR = TensorFormat(
633
+ [b, i, j],
634
+ {b: LevelFormat.DENSE, i: LevelFormat.DENSE, j: LevelFormat.COMPRESSED},
635
+ name="BatchedCSR",
636
+ )
637
+ BatchedDIAINonUniform = TensorFormat(
638
+ [b, i, j],
639
+ {b: LevelFormat.DENSE, j - i: LevelFormat.COMPRESSED, i: LevelFormat.RANGE},
640
+ name="BatchedDIAINonUniform",
641
+ )
642
+ BatchedDIAIUniform = TensorFormat(
643
+ [b, i, j],
644
+ {j - i: LevelFormat.COMPRESSED, b: LevelFormat.DENSE, i: LevelFormat.RANGE},
645
+ name="BatchedDIAIUniform",
646
+ )
647
+
648
+ #
649
+ # Tensor format builders for vectors, matrices, and tensors.
650
+ #
651
+
652
+ @staticmethod
653
+ def BlockVector(blocksize):
654
+ """
655
+ Block vector format.
656
+
657
+ Args:
658
+ blocksize: an integer specifying the blocksize in the i'th dimension.
659
+ """
660
+ i = NamedFormats.i
661
+ name = f"BlockVector{blocksize}"
662
+ return TensorFormat(
663
+ [i],
664
+ {i // blocksize: LevelFormat.COMPRESSED, i % blocksize: LevelFormat.DENSE},
665
+ name=name,
666
+ )
667
+
668
+ @staticmethod
669
+ def BSRRight(blocksize):
670
+ """
671
+ BSRRight matrix format.
672
+
673
+ Args:
674
+ blocksize: a pair specifying the blocksize in the i'th and j'th dimensions.
675
+ """
676
+ i, j = NamedFormats.i, NamedFormats.j
677
+ m, n = blocksize
678
+ name = f"BSRRight{m}x{n}"
679
+ return TensorFormat(
680
+ [i, j],
681
+ {
682
+ i // m: LevelFormat.DENSE,
683
+ j // n: LevelFormat.COMPRESSED,
684
+ i % m: LevelFormat.DENSE,
685
+ j % n: LevelFormat.DENSE,
686
+ },
687
+ name=name,
688
+ )
689
+
690
+ @staticmethod
691
+ def BSRLeft(blocksize):
692
+ """
693
+ BSRLeft matrix format.
694
+
695
+ Args:
696
+ blocksize: a pair specifying the blocksize in the i'th and j'th dimensions.
697
+ """
698
+ i, j = NamedFormats.i, NamedFormats.j
699
+ m, n = blocksize
700
+ name = f"BSRLeft{m}x{n}"
701
+ return TensorFormat(
702
+ [i, j],
703
+ {
704
+ i // m: LevelFormat.DENSE,
705
+ j // n: LevelFormat.COMPRESSED,
706
+ j % n: LevelFormat.DENSE,
707
+ i % m: LevelFormat.DENSE,
708
+ },
709
+ name=name,
710
+ )
711
+
712
+ @staticmethod
713
+ def BSCRight(blocksize):
714
+ """
715
+ BSCRight matrix format.
716
+
717
+ Args:
718
+ blocksize: a pair specifying the blocksize in the i'th and j'th dimensions.
719
+ """
720
+ i, j = NamedFormats.i, NamedFormats.j
721
+ m, n = blocksize
722
+ name = f"BSCRow{m}x{n}"
723
+ return TensorFormat(
724
+ [i, j],
725
+ {
726
+ j // n: LevelFormat.DENSE,
727
+ i // m: LevelFormat.COMPRESSED,
728
+ i % m: LevelFormat.DENSE,
729
+ j % n: LevelFormat.DENSE,
730
+ },
731
+ name=name,
732
+ )
733
+
734
+ @staticmethod
735
+ def BSCLeft(blocksize):
736
+ """
737
+ BSCLeft matrix format.
738
+
739
+ Args:
740
+ blocksize: a pair specifying the blocksize in the i'th and j'th dimensions.
741
+ """
742
+ i, j = NamedFormats.i, NamedFormats.j
743
+ m, n = blocksize
744
+ name = f"BSCCol{m}x{n}"
745
+ return TensorFormat(
746
+ [i, j],
747
+ {
748
+ j // n: LevelFormat.DENSE,
749
+ i // m: LevelFormat.COMPRESSED,
750
+ j % n: LevelFormat.DENSE,
751
+ i % m: LevelFormat.DENSE,
752
+ },
753
+ name=name,
754
+ )
755
+
756
+ @staticmethod
757
+ def DELTA(delta):
758
+ """
759
+ Delta compression matrix format.
760
+
761
+ Args:
762
+ delta: number of bits for delta
763
+ """
764
+ i, j = NamedFormats.i, NamedFormats.j
765
+ name = f"Delta{delta}"
766
+ return TensorFormat(
767
+ [i, j],
768
+ {i: LevelFormat.DENSE, j: (LevelFormat.DELTA, delta)},
769
+ name=name,
770
+ )
771
+
772
+ @staticmethod
773
+ def BSR3(blocksize):
774
+ """
775
+ 3D BSR format.
776
+
777
+ Args:
778
+ blocksize: a triple specifying the blocksize in the i'th, j'th, and k'th
779
+ dimensions.
780
+ """
781
+ i, j, k = NamedFormats.i, NamedFormats.j, NamedFormats.k
782
+ m, n, o = blocksize
783
+ name = f"BSR3{m}x{n}x{o}"
784
+ return TensorFormat(
785
+ [i, j, k],
786
+ {
787
+ i // m: LevelFormat.DENSE,
788
+ j // n: LevelFormat.COMPRESSED,
789
+ k // o: LevelFormat.COMPRESSED,
790
+ i % m: LevelFormat.DENSE,
791
+ j % n: LevelFormat.DENSE,
792
+ k % o: LevelFormat.DENSE,
793
+ },
794
+ name=name,
795
+ )
796
+
797
+ @staticmethod
798
+ def DensedRight(dim):
799
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
800
+ return TensorFormat(
801
+ dims,
802
+ {dims[d]: LevelFormat.DENSE for d in range(dim)},
803
+ name=f"Dense{dim}Right",
804
+ )
805
+
806
+ @staticmethod
807
+ def DensedLeft(dim):
808
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
809
+ return TensorFormat(
810
+ dims,
811
+ {dims[dim - d - 1]: LevelFormat.DENSE for d in range(dim)},
812
+ name=f"Dense{dim}Left",
813
+ )
814
+
815
+ @staticmethod
816
+ def COOd(sparse_dim, dense_dim=0):
817
+ dim = sparse_dim + dense_dim
818
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
819
+ return TensorFormat(
820
+ dims,
821
+ {
822
+ dims[d]: (
823
+ (LevelFormat.COMPRESSED if sparse_dim == 1 else (LevelFormat.COMPRESSED, LevelProperty.NONUNIQUE))
824
+ if d == 0
825
+ else LevelFormat.SINGLETON
826
+ if d < sparse_dim
827
+ else LevelFormat.DENSE
828
+ )
829
+ for d in range(dim)
830
+ },
831
+ name=("COO" if dim == 2 else f"COO{dim}") if dense_dim == 0 else None,
832
+ )
833
+
834
+ @staticmethod
835
+ def CSRd(batch_dim=0, dense_dim=0):
836
+ dim = batch_dim + 2 + dense_dim
837
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
838
+
839
+ name = "CSR"
840
+ if batch_dim > 0:
841
+ name += f"-b{batch_dim}"
842
+ if dense_dim > 0:
843
+ name += f"-d{dense_dim}"
844
+
845
+ return TensorFormat(
846
+ dims,
847
+ {
848
+ dims[d]:
849
+ # Format
850
+ LevelFormat.BATCH if d < batch_dim else LevelFormat.COMPRESSED if d == batch_dim + 1 else LevelFormat.DENSE
851
+ for d in range(dim)
852
+ },
853
+ name=name,
854
+ )
855
+
856
+ @staticmethod
857
+ def CSCd(batch_dim=0, dense_dim=0):
858
+ dim = batch_dim + 2 + dense_dim
859
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
860
+
861
+ name = "CSC"
862
+ if batch_dim > 0:
863
+ name += f"-b{batch_dim}"
864
+ if dense_dim > 0:
865
+ name += f"-d{dense_dim}"
866
+
867
+ return TensorFormat(
868
+ dims,
869
+ {
870
+ dims[d + 1 if d == batch_dim else d - 1 if d == batch_dim + 1 else d]:
871
+ # Format
872
+ LevelFormat.BATCH if d < batch_dim else LevelFormat.COMPRESSED if d == batch_dim + 1 else LevelFormat.DENSE
873
+ for d in range(dim)
874
+ },
875
+ name=name,
876
+ )
877
+
878
+ @staticmethod
879
+ def BSRRightd(blocksize, batch_dim=0, dense_dim=0):
880
+ m, n = blocksize
881
+ dim = batch_dim + 2 + dense_dim
882
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
883
+
884
+ name = f"BSRRight{m}x{n}"
885
+ if batch_dim > 0:
886
+ name += f"-b{batch_dim}"
887
+ if dense_dim > 0:
888
+ name += f"-d{dense_dim}"
889
+
890
+ return TensorFormat(
891
+ dims,
892
+ {dims[d]: LevelFormat.BATCH for d in range(batch_dim)}
893
+ | {
894
+ dims[batch_dim] // m: LevelFormat.DENSE,
895
+ dims[batch_dim + 1] // n: LevelFormat.COMPRESSED,
896
+ dims[batch_dim] % m: LevelFormat.DENSE,
897
+ dims[batch_dim + 1] % n: LevelFormat.DENSE,
898
+ }
899
+ | {dims[d]: LevelFormat.DENSE for d in range(batch_dim + 2, dim)},
900
+ name=name,
901
+ )
902
+
903
+ @staticmethod
904
+ def BSCRightd(blocksize, batch_dim=0, dense_dim=0):
905
+ m, n = blocksize
906
+ dim = batch_dim + 2 + dense_dim
907
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
908
+
909
+ name = f"BSCRight{m}x{n}"
910
+ if batch_dim > 0:
911
+ name += f"-b{batch_dim}"
912
+ if dense_dim > 0:
913
+ name += f"-d{dense_dim}"
914
+
915
+ return TensorFormat(
916
+ dims,
917
+ {dims[d]: LevelFormat.BATCH for d in range(batch_dim)}
918
+ | {
919
+ dims[batch_dim + 1] // n: LevelFormat.DENSE,
920
+ dims[batch_dim] // m: LevelFormat.COMPRESSED,
921
+ dims[batch_dim] % m: LevelFormat.DENSE,
922
+ dims[batch_dim + 1] % n: LevelFormat.DENSE,
923
+ }
924
+ | {dims[d]: LevelFormat.DENSE for d in range(batch_dim + 2, dim)},
925
+ name=name,
926
+ )
927
+
928
+ @staticmethod
929
+ def CSFd(dim):
930
+ dims = [Dimension(dimension_name=f"{chr(ord('i') + d)}") for d in range(dim)]
931
+ return TensorFormat(
932
+ dims,
933
+ {dims[d]: LevelFormat.COMPRESSED for d in range(dim)},
934
+ name=f"CSF{dim}",
935
+ )