cotengrust 0.1.0__tar.gz → 0.1.2__tar.gz

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.
@@ -1,7 +1,7 @@
1
- # This file is autogenerated by maturin v0.15.2
1
+ # This file is autogenerated by maturin v1.2.3
2
2
  # To update, run
3
3
  #
4
- # maturin generate-ci github
4
+ # maturin generate-ci github --pytest
5
5
  #
6
6
  name: CI
7
7
 
@@ -41,6 +41,30 @@ jobs:
41
41
  with:
42
42
  name: wheels
43
43
  path: dist
44
+ - name: pytest
45
+ if: ${{ startsWith(matrix.target, 'x86_64') }}
46
+ shell: bash
47
+ run: |
48
+ set -e
49
+ ls dist/*
50
+ pip install cotengrust --find-links dist --force-reinstall
51
+ pip install pytest numpy cotengra
52
+ pytest --verbose
53
+ - name: pytest
54
+ if: ${{ !startsWith(matrix.target, 'x86') && matrix.target != 'ppc64' }}
55
+ uses: uraimo/run-on-arch-action@v2.5.0
56
+ with:
57
+ arch: ${{ matrix.target }}
58
+ distro: ubuntu22.04
59
+ githubToken: ${{ github.token }}
60
+ install: |
61
+ apt-get update
62
+ apt-get install -y --no-install-recommends python3 python3-pip
63
+ pip3 install -U pip pytest # numpy cotengra
64
+ run: |
65
+ set -e
66
+ pip3 install cotengrust --find-links dist --force-reinstall
67
+ pytest --verbose
44
68
 
45
69
  windows:
46
70
  runs-on: windows-latest
@@ -64,6 +88,15 @@ jobs:
64
88
  with:
65
89
  name: wheels
66
90
  path: dist
91
+ - name: pytest
92
+ if: ${{ !startsWith(matrix.target, 'aarch64') }}
93
+ shell: bash
94
+ run: |
95
+ set -e
96
+ ls dist/*
97
+ pip install cotengrust --find-links dist --force-reinstall
98
+ pip install pytest numpy cotengra
99
+ pytest --verbose
67
100
 
68
101
  macos:
69
102
  runs-on: macos-latest
@@ -86,6 +119,15 @@ jobs:
86
119
  with:
87
120
  name: wheels
88
121
  path: dist
122
+ - name: pytest
123
+ if: ${{ !startsWith(matrix.target, 'aarch64') }}
124
+ shell: bash
125
+ run: |
126
+ set -e
127
+ ls dist/*
128
+ pip install cotengrust --find-links dist --force-reinstall
129
+ pip install pytest numpy cotengra
130
+ pytest --verbose
89
131
 
90
132
  sdist:
91
133
  runs-on: ubuntu-latest
@@ -117,4 +159,4 @@ jobs:
117
159
  MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
118
160
  with:
119
161
  command: upload
120
- args: --skip-existing *
162
+ args: --non-interactive --skip-existing *
@@ -37,7 +37,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
37
37
 
38
38
  [[package]]
39
39
  name = "cotengrust"
40
- version = "0.1.0"
40
+ version = "0.1.2"
41
41
  dependencies = [
42
42
  "bit-set",
43
43
  "ordered-float",
@@ -57,11 +57,17 @@ dependencies = [
57
57
  "wasi",
58
58
  ]
59
59
 
60
+ [[package]]
61
+ name = "heck"
62
+ version = "0.4.1"
63
+ source = "registry+https://github.com/rust-lang/crates.io-index"
64
+ checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
65
+
60
66
  [[package]]
61
67
  name = "indoc"
62
- version = "1.0.9"
68
+ version = "2.0.4"
63
69
  source = "registry+https://github.com/rust-lang/crates.io-index"
64
- checksum = "bfa799dd5ed20a7e349f3b4639aa80d74549c81716d9ec4f994c9b5815598306"
70
+ checksum = "1e186cfbae8084e513daff4240b4797e342f988cecda4fb6c939150f96315fd8"
65
71
 
66
72
  [[package]]
67
73
  name = "libc"
@@ -105,9 +111,9 @@ checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
105
111
 
106
112
  [[package]]
107
113
  name = "ordered-float"
108
- version = "3.9.1"
114
+ version = "4.2.0"
109
115
  source = "registry+https://github.com/rust-lang/crates.io-index"
110
- checksum = "2a54938017eacd63036332b4ae5c8a49fc8c0c1d6d629893057e4f13609edd06"
116
+ checksum = "a76df7075c7d4d01fdcb46c912dd17fba5b60c78ea480b475f2b6ab6f666584e"
111
117
  dependencies = [
112
118
  "num-traits",
113
119
  ]
@@ -135,6 +141,12 @@ dependencies = [
135
141
  "windows-targets",
136
142
  ]
137
143
 
144
+ [[package]]
145
+ name = "portable-atomic"
146
+ version = "1.6.0"
147
+ source = "registry+https://github.com/rust-lang/crates.io-index"
148
+ checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0"
149
+
138
150
  [[package]]
139
151
  name = "ppv-lite86"
140
152
  version = "0.2.17"
@@ -152,15 +164,16 @@ dependencies = [
152
164
 
153
165
  [[package]]
154
166
  name = "pyo3"
155
- version = "0.19.2"
167
+ version = "0.21.1"
156
168
  source = "registry+https://github.com/rust-lang/crates.io-index"
157
- checksum = "e681a6cfdc4adcc93b4d3cf993749a4552018ee0a9b65fc0ccfad74352c72a38"
169
+ checksum = "a7a8b1990bd018761768d5e608a13df8bd1ac5f678456e0f301bb93e5f3ea16b"
158
170
  dependencies = [
159
171
  "cfg-if",
160
172
  "indoc",
161
173
  "libc",
162
174
  "memoffset",
163
175
  "parking_lot",
176
+ "portable-atomic",
164
177
  "pyo3-build-config",
165
178
  "pyo3-ffi",
166
179
  "pyo3-macros",
@@ -169,9 +182,9 @@ dependencies = [
169
182
 
170
183
  [[package]]
171
184
  name = "pyo3-build-config"
172
- version = "0.19.2"
185
+ version = "0.21.1"
173
186
  source = "registry+https://github.com/rust-lang/crates.io-index"
174
- checksum = "076c73d0bc438f7a4ef6fdd0c3bb4732149136abd952b110ac93e4edb13a6ba5"
187
+ checksum = "650dca34d463b6cdbdb02b1d71bfd6eb6b6816afc708faebb3bac1380ff4aef7"
175
188
  dependencies = [
176
189
  "once_cell",
177
190
  "target-lexicon",
@@ -179,9 +192,9 @@ dependencies = [
179
192
 
180
193
  [[package]]
181
194
  name = "pyo3-ffi"
182
- version = "0.19.2"
195
+ version = "0.21.1"
183
196
  source = "registry+https://github.com/rust-lang/crates.io-index"
184
- checksum = "e53cee42e77ebe256066ba8aa77eff722b3bb91f3419177cf4cd0f304d3284d9"
197
+ checksum = "09a7da8fc04a8a2084909b59f29e1b8474decac98b951d77b80b26dc45f046ad"
185
198
  dependencies = [
186
199
  "libc",
187
200
  "pyo3-build-config",
@@ -189,9 +202,9 @@ dependencies = [
189
202
 
190
203
  [[package]]
191
204
  name = "pyo3-macros"
192
- version = "0.19.2"
205
+ version = "0.21.1"
193
206
  source = "registry+https://github.com/rust-lang/crates.io-index"
194
- checksum = "dfeb4c99597e136528c6dd7d5e3de5434d1ceaf487436a3f03b2d56b6fc9efd1"
207
+ checksum = "4b8a199fce11ebb28e3569387228836ea98110e43a804a530a9fd83ade36d513"
195
208
  dependencies = [
196
209
  "proc-macro2",
197
210
  "pyo3-macros-backend",
@@ -201,11 +214,13 @@ dependencies = [
201
214
 
202
215
  [[package]]
203
216
  name = "pyo3-macros-backend"
204
- version = "0.19.2"
217
+ version = "0.21.1"
205
218
  source = "registry+https://github.com/rust-lang/crates.io-index"
206
- checksum = "947dc12175c254889edc0c02e399476c2f652b4b9ebd123aa655c224de259536"
219
+ checksum = "93fbbfd7eb553d10036513cb122b888dcd362a945a00b06c165f2ab480d4cc3b"
207
220
  dependencies = [
221
+ "heck",
208
222
  "proc-macro2",
223
+ "pyo3-build-config",
209
224
  "quote",
210
225
  "syn",
211
226
  ]
@@ -278,9 +293,9 @@ checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9"
278
293
 
279
294
  [[package]]
280
295
  name = "syn"
281
- version = "1.0.109"
296
+ version = "2.0.32"
282
297
  source = "registry+https://github.com/rust-lang/crates.io-index"
283
- checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
298
+ checksum = "239814284fd6f1a4ffe4ca893952cdd93c224b6a1571c9a9eadd670295c0c9e2"
284
299
  dependencies = [
285
300
  "proc-macro2",
286
301
  "quote",
@@ -301,9 +316,9 @@ checksum = "301abaae475aa91687eb82514b328ab47a211a533026cb25fc3e519b86adfc3c"
301
316
 
302
317
  [[package]]
303
318
  name = "unindent"
304
- version = "0.1.11"
319
+ version = "0.2.3"
305
320
  source = "registry+https://github.com/rust-lang/crates.io-index"
306
- checksum = "e1766d682d402817b5ac4490b3c3002d91dfa0d22812f341609f97b08757359c"
321
+ checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce"
307
322
 
308
323
  [[package]]
309
324
  name = "wasi"
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "cotengrust"
3
- version = "0.1.0"
3
+ version = "0.1.2"
4
4
  edition = "2021"
5
5
 
6
6
  # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -10,8 +10,8 @@ crate-type = ["cdylib"]
10
10
 
11
11
  [dependencies]
12
12
  bit-set = "0.5"
13
- pyo3 = "0.19"
14
- ordered-float = "3.9"
13
+ ordered-float = "4.2"
14
+ pyo3 = "0.21"
15
15
  rand = "0.8"
16
16
  rustc-hash = "1.1"
17
17
 
@@ -19,4 +19,3 @@ rustc-hash = "1.1"
19
19
  codegen-units = 1
20
20
  lto = true
21
21
  opt-level = 3
22
- panic = "abort"
@@ -0,0 +1,249 @@
1
+ Metadata-Version: 2.3
2
+ Name: cotengrust
3
+ Version: 0.1.2
4
+ Classifier: Programming Language :: Rust
5
+ Classifier: Programming Language :: Python :: Implementation :: CPython
6
+ Classifier: Programming Language :: Python :: Implementation :: PyPy
7
+ License-File: LICENSE
8
+ Summary: Fast contraction ordering primitives for tensor networks.
9
+ Author-email: Johnnie Gray <johnniemcgray@gmail.com>
10
+ Requires-Python: >=3.8
11
+ Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
12
+
13
+ # cotengrust
14
+
15
+ `cotengrust` provides fast rust implementations of contraction ordering
16
+ primitives for tensor networks or einsum expressions. The two main functions
17
+ are:
18
+
19
+ - `optimize_optimal(inputs, output, size_dict, **kwargs)`
20
+ - `optimize_greedy(inputs, output, size_dict, **kwargs)`
21
+
22
+ The optimal algorithm is an optimized version of the `opt_einsum` 'dp'
23
+ path - itself an implementation of https://arxiv.org/abs/1304.6112.
24
+
25
+
26
+ ## Installation
27
+
28
+ `cotengrust` is available for most platforms from
29
+ [PyPI](https://pypi.org/project/cotengrust/):
30
+
31
+ ```bash
32
+ pip install cotengrust
33
+ ```
34
+
35
+ or if you want to develop locally (which requires [pyo3](https://github.com/PyO3/pyo3)
36
+ and [maturin](https://github.com/PyO3/maturin)):
37
+
38
+ ```bash
39
+ git clone https://github.com/jcmgray/cotengrust.git
40
+ cd cotengrust
41
+ maturin develop --release
42
+ ```
43
+ (the release flag is very important for assessing performance!).
44
+
45
+
46
+ ## Usage
47
+
48
+ If `cotengrust` is installed, then by default `cotengra` will use it for its
49
+ greedy and optimal subroutines, notably subtree reconfiguration. You can also
50
+ call the routines directly:
51
+
52
+ ```python
53
+ import cotengra as ctg
54
+ import cotengrust as ctgr
55
+
56
+ # specify an 8x8 square lattice contraction
57
+ inputs, output, shapes, size_dict = ctg.utils.lattice_equation([8, 8])
58
+
59
+ # find the optimal 'combo' contraction path
60
+ %%time
61
+ path = ctgr.optimize_optimal(inputs, output, size_dict, minimize='combo')
62
+ # CPU times: user 13.7 s, sys: 83.4 ms, total: 13.7 s
63
+ # Wall time: 13.7 s
64
+
65
+ # construct a contraction tree for further introspection
66
+ tree = ctg.ContractionTree.from_path(
67
+ inputs, output, size_dict, path=path
68
+ )
69
+ tree.plot_rubberband()
70
+ ```
71
+ ![optimal-8x8-order](https://github.com/jcmgray/cotengrust/assets/8982598/f8e18ff2-5ace-4e46-81e1-06bffaef5e45)
72
+
73
+
74
+ ## API
75
+
76
+ The optimize functions follow the api of the python implementations in `cotengra.pathfinders.path_basic.py`.
77
+
78
+ ```python
79
+ def optimize_optimal(
80
+ inputs,
81
+ output,
82
+ size_dict,
83
+ minimize='flops',
84
+ cost_cap=2,
85
+ search_outer=False,
86
+ simplify=True,
87
+ use_ssa=False,
88
+ ):
89
+ """Find an optimal contraction ordering.
90
+
91
+ Parameters
92
+ ----------
93
+ inputs : Sequence[Sequence[str]]
94
+ The indices of each input tensor.
95
+ output : Sequence[str]
96
+ The indices of the output tensor.
97
+ size_dict : dict[str, int]
98
+ The size of each index.
99
+ minimize : str, optional
100
+ The cost function to minimize. The options are:
101
+
102
+ - "flops": minimize with respect to total operation count only
103
+ (also known as contraction cost)
104
+ - "size": minimize with respect to maximum intermediate size only
105
+ (also known as contraction width)
106
+ - 'write' : minimize the sum of all tensor sizes, i.e. memory written
107
+ - 'combo' or 'combo={factor}` : minimize the sum of
108
+ FLOPS + factor * WRITE, with a default factor of 64.
109
+ - 'limit' or 'limit={factor}` : minimize the sum of
110
+ MAX(FLOPS, alpha * WRITE) for each individual contraction, with a
111
+ default factor of 64.
112
+
113
+ 'combo' is generally a good default in term of practical hardware
114
+ performance, where both memory bandwidth and compute are limited.
115
+ cost_cap : float, optional
116
+ The maximum cost of a contraction to initially consider. This acts like
117
+ a sieve and is doubled at each iteration until the optimal path can
118
+ be found, but supplying an accurate guess can speed up the algorithm.
119
+ search_outer : bool, optional
120
+ If True, consider outer product contractions. This is much slower but
121
+ theoretically might be required to find the true optimal 'flops'
122
+ ordering. In practical settings (i.e. with minimize='combo'), outer
123
+ products should not be required.
124
+ simplify : bool, optional
125
+ Whether to perform simplifications before optimizing. These are:
126
+
127
+ - ignore any indices that appear in all terms
128
+ - combine any repeated indices within a single term
129
+ - reduce any non-output indices that only appear on a single term
130
+ - combine any scalar terms
131
+ - combine any tensors with matching indices (hadamard products)
132
+
133
+ Such simpifications may be required in the general case for the proper
134
+ functioning of the core optimization, but may be skipped if the input
135
+ indices are already in a simplified form.
136
+ use_ssa : bool, optional
137
+ Whether to return the contraction path in 'single static assignment'
138
+ (SSA) format (i.e. as if each intermediate is appended to the list of
139
+ inputs, without removals). This can be quicker and easier to work with
140
+ than the 'linear recycled' format that `numpy` and `opt_einsum` use.
141
+
142
+ Returns
143
+ -------
144
+ path : list[list[int]]
145
+ The contraction path, given as a sequence of pairs of node indices. It
146
+ may also have single term contractions if `simplify=True`.
147
+ """
148
+ ...
149
+
150
+
151
+ def optimize_greedy(
152
+ inputs,
153
+ output,
154
+ size_dict,
155
+ costmod=1.0,
156
+ temperature=0.0,
157
+ simplify=True,
158
+ use_ssa=False,
159
+ ):
160
+ """Find a contraction path using a (randomizable) greedy algorithm.
161
+
162
+ Parameters
163
+ ----------
164
+ inputs : Sequence[Sequence[str]]
165
+ The indices of each input tensor.
166
+ output : Sequence[str]
167
+ The indices of the output tensor.
168
+ size_dict : dict[str, int]
169
+ A dictionary mapping indices to their dimension.
170
+ costmod : float, optional
171
+ When assessing local greedy scores how much to weight the size of the
172
+ tensors removed compared to the size of the tensor added::
173
+
174
+ score = size_ab - costmod * (size_a + size_b)
175
+
176
+ This can be a useful hyper-parameter to tune.
177
+ temperature : float, optional
178
+ When asessing local greedy scores, how much to randomly perturb the
179
+ score. This is implemented as::
180
+
181
+ score -> sign(score) * log(|score|) - temperature * gumbel()
182
+
183
+ which implements boltzmann sampling.
184
+ simplify : bool, optional
185
+ Whether to perform simplifications before optimizing. These are:
186
+
187
+ - ignore any indices that appear in all terms
188
+ - combine any repeated indices within a single term
189
+ - reduce any non-output indices that only appear on a single term
190
+ - combine any scalar terms
191
+ - combine any tensors with matching indices (hadamard products)
192
+
193
+ Such simpifications may be required in the general case for the proper
194
+ functioning of the core optimization, but may be skipped if the input
195
+ indices are already in a simplified form.
196
+ use_ssa : bool, optional
197
+ Whether to return the contraction path in 'single static assignment'
198
+ (SSA) format (i.e. as if each intermediate is appended to the list of
199
+ inputs, without removals). This can be quicker and easier to work with
200
+ than the 'linear recycled' format that `numpy` and `opt_einsum` use.
201
+
202
+ Returns
203
+ -------
204
+ path : list[list[int]]
205
+ The contraction path, given as a sequence of pairs of node indices. It
206
+ may also have single term contractions if `simplify=True`.
207
+ """
208
+
209
+ def optimize_simplify(
210
+ inputs,
211
+ output,
212
+ size_dict,
213
+ use_ssa=False,
214
+ ):
215
+ """Find the (partial) contracton path for simplifiactions only.
216
+
217
+ Parameters
218
+ ----------
219
+ inputs : Sequence[Sequence[str]]
220
+ The indices of each input tensor.
221
+ output : Sequence[str]
222
+ The indices of the output tensor.
223
+ size_dict : dict[str, int]
224
+ A dictionary mapping indices to their dimension.
225
+ use_ssa : bool, optional
226
+ Whether to return the contraction path in 'single static assignment'
227
+ (SSA) format (i.e. as if each intermediate is appended to the list of
228
+ inputs, without removals). This can be quicker and easier to work with
229
+ than the 'linear recycled' format that `numpy` and `opt_einsum` use.
230
+
231
+ Returns
232
+ -------
233
+ path : list[list[int]]
234
+ The contraction path, given as a sequence of pairs of node indices. It
235
+ may also have single term contractions.
236
+
237
+ """
238
+ ...
239
+
240
+ def ssa_to_linear(ssa_path, n=None):
241
+ """Convert a SSA path to linear format."""
242
+ ...
243
+
244
+ def find_subgraphs(inputs, output, size_dict,):
245
+ """Find all disconnected subgraphs of a specified contraction."""
246
+ ...
247
+ ```
248
+
249
+