rigetti-quax 0.2.1__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.
@@ -0,0 +1,11 @@
1
+ Licensed under the Apache License, Version 2.0 (the "License");
2
+ you may not use this file except in compliance with the License.
3
+ You may obtain a copy of the License at
4
+
5
+ www.apache.org
6
+
7
+ Unless required by applicable law or agreed to in writing, software
8
+ distributed under the License is distributed on an "AS IS" BASIS,
9
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10
+ See the License for the specific language governing permissions and
11
+ limitations under the License.
@@ -0,0 +1,171 @@
1
+ Metadata-Version: 2.4
2
+ Name: rigetti-quax
3
+ Version: 0.2.1
4
+ Summary: A high-performance library for quantum information science built on top of JAX
5
+ License-Expression: Apache-2.0
6
+ License-File: LICENSE
7
+ Author: Bram Evert
8
+ Author-email: bevert@rigetti.com
9
+ Requires-Python: >=3.12
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
13
+ Classifier: Programming Language :: Python :: 3.14
14
+ Requires-Dist: jax (>=0.8.2)
15
+ Description-Content-Type: text/markdown
16
+
17
+ # Quax
18
+
19
+ A high-performance quantum information science library built on top of JAX.
20
+
21
+ ## Data Representation
22
+
23
+ Quax stores all quantum objects in **tensor format**, preserving the structure of individual qudits for efficient tensor network operations. The `.matrix` property provides the flattened matrix representation when needed.
24
+
25
+ | Type | Tensor Shape | Matrix Shape |
26
+ |------|--------------|--------------|
27
+ | StateVector | `(*ensemble, d0, d1, ...)` | `(*ensemble, prod(dims))` |
28
+ | DensityMatrix | `(*ensemble, d0_out, ..., d0_in, ...)` | `(*ensemble, prod(dims), prod(dims))` |
29
+ | Unitary/Operator | `(*ensemble, d0_out, ..., d0_in, ...)` | `(*ensemble, prod(dims_out), prod(dims_in))` |
30
+ | KrausMap | `(*ensemble, num_kraus, d0_out, ..., d0_in, ...)` | `(*ensemble, num_kraus, d_out, d_in)` |
31
+ | SuperOp/Choi/PauliLiouville | `(*ensemble, d_out_bra..., d_out_ket..., d_in_bra..., d_in_ket...)` | `(*ensemble, prod(dims_out)², prod(dims_in)²)` |
32
+
33
+ ## Supported Operations on Quantum Objects
34
+
35
+ ### Unary Operations
36
+
37
+ | Operation | StateVector | DensityMatrix | Unitary | Kraus | SuperOp | KrausMap | Choi | Chi | PauliLiouville |
38
+ |-----------|-------------|---------------|---------|-------|---------|----------|------|-----|----------------|
39
+ | `-x` (negation) | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
40
+ | `x * scalar` | ✓ | ✓ | ✓¹ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
41
+ | `x.conj()` | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
42
+ | `x.T` (transpose) | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
43
+ | `x.h` (hermitian) | ✓³ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
44
+ | `x ** n` (power) | ✗ | ✓ | ✓ | ✗ | ✓ | ✓ | ✓ | ✗ | ✓ |
45
+
46
+ ¹ Returns `Unitary` if `|scalar| = 1`, otherwise `Kraus`
47
+ ² Returns self (no-op for vectors)
48
+ ³ Returns `conj()` for vectors
49
+
50
+ ---
51
+
52
+
53
+ ### Binary Operations: Composition (`@`)
54
+
55
+ | Operation | Supported | Output Type |
56
+ |-----------|-----------|-------------|
57
+ | `StateVector @ StateVector` | ✓ | scalar |
58
+ | `StateVector @ DensityMatrix` | ✓ | StateVector |
59
+ | `StateVector @ Unitary` | ✗ | — |
60
+ | `StateVector @ SuperOp` | ✗ | — |
61
+ | `StateVector @ KrausMap` | ✗ | — |
62
+ | `StateVector @ Choi` | ✗ | — |
63
+ | `StateVector @ PauliLiouville` | ✗ | — |
64
+ | `DensityMatrix @ StateVector` | ✓ | StateVector |
65
+ | `DensityMatrix @ DensityMatrix` | ✓ | DensityMatrix |
66
+ | `DensityMatrix @ Unitary` | ✗ | — |
67
+ | `DensityMatrix @ SuperOp` | ✗ | — |
68
+ | `DensityMatrix @ KrausMap` | ✗ | — |
69
+ | `DensityMatrix @ Choi` | ✗ | — |
70
+ | `DensityMatrix @ PauliLiouville` | ✗ | — |
71
+ | `Unitary @ StateVector` | ✓ | StateVector |
72
+ | `Unitary @ DensityMatrix` | ✓ | DensityMatrix |
73
+ | `Unitary @ Unitary` | ✓ | Unitary |
74
+ | `Unitary @ SuperOp` | ✓ | SuperOp |
75
+ | `Unitary @ KrausMap` | ✓ | KrausMap |
76
+ | `Unitary @ Choi` | ✓ | Choi |
77
+ | `Unitary @ PauliLiouville` | ✓ | PauliLiouville |
78
+ | `SuperOp @ StateVector` | ✓ | DensityMatrix |
79
+ | `SuperOp @ DensityMatrix` | ✓ | DensityMatrix |
80
+ | `SuperOp @ Unitary` | ✓ | SuperOp |
81
+ | `SuperOp @ SuperOp` | ✓ | SuperOp |
82
+ | `SuperOp @ KrausMap` | ✓ | KrausMap |
83
+ | `SuperOp @ Choi` | ✓ | Choi |
84
+ | `SuperOp @ PauliLiouville` | ✓ | PauliLiouville |
85
+ | `KrausMap @ StateVector` | ✓ | DensityMatrix |
86
+ | `KrausMap @ DensityMatrix` | ✓ | DensityMatrix |
87
+ | `KrausMap @ Unitary` | ✓ | KrausMap |
88
+ | `KrausMap @ SuperOp` | ✓ | SuperOp |
89
+ | `KrausMap @ KrausMap` | ✓ | KrausMap |
90
+ | `KrausMap @ Choi` | ✓ | Choi |
91
+ | `KrausMap @ PauliLiouville` | ✓ | PauliLiouville |
92
+ | `Choi @ StateVector` | ✓ | DensityMatrix |
93
+ | `Choi @ DensityMatrix` | ✓ | DensityMatrix |
94
+ | `Choi @ Unitary` | ✓ | Choi |
95
+ | `Choi @ SuperOp` | ✓ | SuperOp |
96
+ | `Choi @ KrausMap` | ✓ | KrausMap |
97
+ | `Choi @ Choi` | ✓ | Choi |
98
+ | `Choi @ PauliLiouville` | ✓ | PauliLiouville |
99
+ | `PauliLiouville @ StateVector` | ✓ | DensityMatrix |
100
+ | `PauliLiouville @ DensityMatrix` | ✓ | DensityMatrix |
101
+ | `PauliLiouville @ Unitary` | ✓ | PauliLiouville |
102
+ | `PauliLiouville @ SuperOp` | ✓ | SuperOp |
103
+ | `PauliLiouville @ KrausMap` | ✓ | KrausMap |
104
+ | `PauliLiouville @ Choi` | ✓ | Choi |
105
+ | `PauliLiouville @ PauliLiouville` | ✓ | PauliLiouville |
106
+
107
+ ---
108
+
109
+ ### Binary Operations: Tensor Product (`|`)
110
+
111
+ | Operation | Supported | Output Type |
112
+ |-----------|-----------|-------------|
113
+ | `StateVector \| StateVector` | ✓ | StateVector |
114
+ | `StateVector \| DensityMatrix` | ✓ | DensityMatrix |
115
+ | `StateVector \| Unitary` | ✗ | — |
116
+ | `StateVector \| SuperOp` | ✗ | — |
117
+ | `StateVector \| KrausMap` | ✗ | — |
118
+ | `StateVector \| Choi` | ✗ | — |
119
+ | `StateVector \| PauliLiouville` | ✗ | — |
120
+ | `DensityMatrix \| StateVector` | ✓ | DensityMatrix |
121
+ | `DensityMatrix \| DensityMatrix` | ✓ | DensityMatrix |
122
+ | `DensityMatrix \| Unitary` | ✗ | — |
123
+ | `DensityMatrix \| SuperOp` | ✗ | — |
124
+ | `DensityMatrix \| KrausMap` | ✗ | — |
125
+ | `DensityMatrix \| Choi` | ✗ | — |
126
+ | `DensityMatrix \| PauliLiouville` | ✗ | — |
127
+ | `Unitary \| StateVector` | ✗ | — |
128
+ | `Unitary \| DensityMatrix` | ✗ | — |
129
+ | `Unitary \| Unitary` | ✓ | Unitary |
130
+ | `Unitary \| SuperOp` | ✓ | SuperOp |
131
+ | `Unitary \| KrausMap` | ✓ | KrausMap |
132
+ | `Unitary \| Choi` | ✓ | Choi |
133
+ | `Unitary \| PauliLiouville` | ✓ | PauliLiouville |
134
+ | `SuperOp \| StateVector` | ✗ | — |
135
+ | `SuperOp \| DensityMatrix` | ✗ | — |
136
+ | `SuperOp \| Unitary` | ✓ | SuperOp |
137
+ | `SuperOp \| SuperOp` | ✓ | SuperOp |
138
+ | `SuperOp \| KrausMap` | ✓ | KrausMap |
139
+ | `SuperOp \| Choi` | ✓ | Choi |
140
+ | `SuperOp \| PauliLiouville` | ✓ | PauliLiouville |
141
+ | `KrausMap \| StateVector` | ✗ | — |
142
+ | `KrausMap \| DensityMatrix` | ✗ | — |
143
+ | `KrausMap \| Unitary` | ✓ | KrausMap |
144
+ | `KrausMap \| SuperOp` | ✓ | SuperOp |
145
+ | `KrausMap \| KrausMap` | ✓ | KrausMap |
146
+ | `KrausMap \| Choi` | ✓ | Choi |
147
+ | `KrausMap \| PauliLiouville` | ✓ | PauliLiouville |
148
+ | `Choi \| StateVector` | ✗ | — |
149
+ | `Choi \| DensityMatrix` | ✗ | — |
150
+ | `Choi \| Unitary` | ✓ | Choi |
151
+ | `Choi \| SuperOp` | ✓ | SuperOp |
152
+ | `Choi \| KrausMap` | ✓ | KrausMap |
153
+ | `Choi \| Choi` | ✓ | Choi |
154
+ | `Choi \| PauliLiouville` | ✓ | PauliLiouville |
155
+ | `PauliLiouville \| StateVector` | ✗ | — |
156
+ | `PauliLiouville \| DensityMatrix` | ✗ | — |
157
+ | `PauliLiouville \| Unitary` | ✓ | PauliLiouville |
158
+ | `PauliLiouville \| SuperOp` | ✓ | SuperOp |
159
+ | `PauliLiouville \| KrausMap` | ✓ | KrausMap |
160
+ | `PauliLiouville \| Choi` | ✓ | Choi |
161
+ | `PauliLiouville \| PauliLiouville` | ✓ | PauliLiouville |
162
+
163
+ ---
164
+
165
+ ### Notes
166
+
167
+ - **Chi** is not included in binary operations because it has no implemented transformations to/from other representations
168
+ - The composition rules follow the principle that when mixing representations, the result uses the "right" operand's representation type
169
+ - State-Operator tensor products (`State | Operator`) return `NotImplemented`
170
+ - Operator-State tensor products (`Operator | State`) return `NotImplemented`
171
+
@@ -0,0 +1,154 @@
1
+ # Quax
2
+
3
+ A high-performance quantum information science library built on top of JAX.
4
+
5
+ ## Data Representation
6
+
7
+ Quax stores all quantum objects in **tensor format**, preserving the structure of individual qudits for efficient tensor network operations. The `.matrix` property provides the flattened matrix representation when needed.
8
+
9
+ | Type | Tensor Shape | Matrix Shape |
10
+ |------|--------------|--------------|
11
+ | StateVector | `(*ensemble, d0, d1, ...)` | `(*ensemble, prod(dims))` |
12
+ | DensityMatrix | `(*ensemble, d0_out, ..., d0_in, ...)` | `(*ensemble, prod(dims), prod(dims))` |
13
+ | Unitary/Operator | `(*ensemble, d0_out, ..., d0_in, ...)` | `(*ensemble, prod(dims_out), prod(dims_in))` |
14
+ | KrausMap | `(*ensemble, num_kraus, d0_out, ..., d0_in, ...)` | `(*ensemble, num_kraus, d_out, d_in)` |
15
+ | SuperOp/Choi/PauliLiouville | `(*ensemble, d_out_bra..., d_out_ket..., d_in_bra..., d_in_ket...)` | `(*ensemble, prod(dims_out)², prod(dims_in)²)` |
16
+
17
+ ## Supported Operations on Quantum Objects
18
+
19
+ ### Unary Operations
20
+
21
+ | Operation | StateVector | DensityMatrix | Unitary | Kraus | SuperOp | KrausMap | Choi | Chi | PauliLiouville |
22
+ |-----------|-------------|---------------|---------|-------|---------|----------|------|-----|----------------|
23
+ | `-x` (negation) | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
24
+ | `x * scalar` | ✓ | ✓ | ✓¹ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
25
+ | `x.conj()` | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
26
+ | `x.T` (transpose) | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
27
+ | `x.h` (hermitian) | ✓³ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
28
+ | `x ** n` (power) | ✗ | ✓ | ✓ | ✗ | ✓ | ✓ | ✓ | ✗ | ✓ |
29
+
30
+ ¹ Returns `Unitary` if `|scalar| = 1`, otherwise `Kraus`
31
+ ² Returns self (no-op for vectors)
32
+ ³ Returns `conj()` for vectors
33
+
34
+ ---
35
+
36
+
37
+ ### Binary Operations: Composition (`@`)
38
+
39
+ | Operation | Supported | Output Type |
40
+ |-----------|-----------|-------------|
41
+ | `StateVector @ StateVector` | ✓ | scalar |
42
+ | `StateVector @ DensityMatrix` | ✓ | StateVector |
43
+ | `StateVector @ Unitary` | ✗ | — |
44
+ | `StateVector @ SuperOp` | ✗ | — |
45
+ | `StateVector @ KrausMap` | ✗ | — |
46
+ | `StateVector @ Choi` | ✗ | — |
47
+ | `StateVector @ PauliLiouville` | ✗ | — |
48
+ | `DensityMatrix @ StateVector` | ✓ | StateVector |
49
+ | `DensityMatrix @ DensityMatrix` | ✓ | DensityMatrix |
50
+ | `DensityMatrix @ Unitary` | ✗ | — |
51
+ | `DensityMatrix @ SuperOp` | ✗ | — |
52
+ | `DensityMatrix @ KrausMap` | ✗ | — |
53
+ | `DensityMatrix @ Choi` | ✗ | — |
54
+ | `DensityMatrix @ PauliLiouville` | ✗ | — |
55
+ | `Unitary @ StateVector` | ✓ | StateVector |
56
+ | `Unitary @ DensityMatrix` | ✓ | DensityMatrix |
57
+ | `Unitary @ Unitary` | ✓ | Unitary |
58
+ | `Unitary @ SuperOp` | ✓ | SuperOp |
59
+ | `Unitary @ KrausMap` | ✓ | KrausMap |
60
+ | `Unitary @ Choi` | ✓ | Choi |
61
+ | `Unitary @ PauliLiouville` | ✓ | PauliLiouville |
62
+ | `SuperOp @ StateVector` | ✓ | DensityMatrix |
63
+ | `SuperOp @ DensityMatrix` | ✓ | DensityMatrix |
64
+ | `SuperOp @ Unitary` | ✓ | SuperOp |
65
+ | `SuperOp @ SuperOp` | ✓ | SuperOp |
66
+ | `SuperOp @ KrausMap` | ✓ | KrausMap |
67
+ | `SuperOp @ Choi` | ✓ | Choi |
68
+ | `SuperOp @ PauliLiouville` | ✓ | PauliLiouville |
69
+ | `KrausMap @ StateVector` | ✓ | DensityMatrix |
70
+ | `KrausMap @ DensityMatrix` | ✓ | DensityMatrix |
71
+ | `KrausMap @ Unitary` | ✓ | KrausMap |
72
+ | `KrausMap @ SuperOp` | ✓ | SuperOp |
73
+ | `KrausMap @ KrausMap` | ✓ | KrausMap |
74
+ | `KrausMap @ Choi` | ✓ | Choi |
75
+ | `KrausMap @ PauliLiouville` | ✓ | PauliLiouville |
76
+ | `Choi @ StateVector` | ✓ | DensityMatrix |
77
+ | `Choi @ DensityMatrix` | ✓ | DensityMatrix |
78
+ | `Choi @ Unitary` | ✓ | Choi |
79
+ | `Choi @ SuperOp` | ✓ | SuperOp |
80
+ | `Choi @ KrausMap` | ✓ | KrausMap |
81
+ | `Choi @ Choi` | ✓ | Choi |
82
+ | `Choi @ PauliLiouville` | ✓ | PauliLiouville |
83
+ | `PauliLiouville @ StateVector` | ✓ | DensityMatrix |
84
+ | `PauliLiouville @ DensityMatrix` | ✓ | DensityMatrix |
85
+ | `PauliLiouville @ Unitary` | ✓ | PauliLiouville |
86
+ | `PauliLiouville @ SuperOp` | ✓ | SuperOp |
87
+ | `PauliLiouville @ KrausMap` | ✓ | KrausMap |
88
+ | `PauliLiouville @ Choi` | ✓ | Choi |
89
+ | `PauliLiouville @ PauliLiouville` | ✓ | PauliLiouville |
90
+
91
+ ---
92
+
93
+ ### Binary Operations: Tensor Product (`|`)
94
+
95
+ | Operation | Supported | Output Type |
96
+ |-----------|-----------|-------------|
97
+ | `StateVector \| StateVector` | ✓ | StateVector |
98
+ | `StateVector \| DensityMatrix` | ✓ | DensityMatrix |
99
+ | `StateVector \| Unitary` | ✗ | — |
100
+ | `StateVector \| SuperOp` | ✗ | — |
101
+ | `StateVector \| KrausMap` | ✗ | — |
102
+ | `StateVector \| Choi` | ✗ | — |
103
+ | `StateVector \| PauliLiouville` | ✗ | — |
104
+ | `DensityMatrix \| StateVector` | ✓ | DensityMatrix |
105
+ | `DensityMatrix \| DensityMatrix` | ✓ | DensityMatrix |
106
+ | `DensityMatrix \| Unitary` | ✗ | — |
107
+ | `DensityMatrix \| SuperOp` | ✗ | — |
108
+ | `DensityMatrix \| KrausMap` | ✗ | — |
109
+ | `DensityMatrix \| Choi` | ✗ | — |
110
+ | `DensityMatrix \| PauliLiouville` | ✗ | — |
111
+ | `Unitary \| StateVector` | ✗ | — |
112
+ | `Unitary \| DensityMatrix` | ✗ | — |
113
+ | `Unitary \| Unitary` | ✓ | Unitary |
114
+ | `Unitary \| SuperOp` | ✓ | SuperOp |
115
+ | `Unitary \| KrausMap` | ✓ | KrausMap |
116
+ | `Unitary \| Choi` | ✓ | Choi |
117
+ | `Unitary \| PauliLiouville` | ✓ | PauliLiouville |
118
+ | `SuperOp \| StateVector` | ✗ | — |
119
+ | `SuperOp \| DensityMatrix` | ✗ | — |
120
+ | `SuperOp \| Unitary` | ✓ | SuperOp |
121
+ | `SuperOp \| SuperOp` | ✓ | SuperOp |
122
+ | `SuperOp \| KrausMap` | ✓ | KrausMap |
123
+ | `SuperOp \| Choi` | ✓ | Choi |
124
+ | `SuperOp \| PauliLiouville` | ✓ | PauliLiouville |
125
+ | `KrausMap \| StateVector` | ✗ | — |
126
+ | `KrausMap \| DensityMatrix` | ✗ | — |
127
+ | `KrausMap \| Unitary` | ✓ | KrausMap |
128
+ | `KrausMap \| SuperOp` | ✓ | SuperOp |
129
+ | `KrausMap \| KrausMap` | ✓ | KrausMap |
130
+ | `KrausMap \| Choi` | ✓ | Choi |
131
+ | `KrausMap \| PauliLiouville` | ✓ | PauliLiouville |
132
+ | `Choi \| StateVector` | ✗ | — |
133
+ | `Choi \| DensityMatrix` | ✗ | — |
134
+ | `Choi \| Unitary` | ✓ | Choi |
135
+ | `Choi \| SuperOp` | ✓ | SuperOp |
136
+ | `Choi \| KrausMap` | ✓ | KrausMap |
137
+ | `Choi \| Choi` | ✓ | Choi |
138
+ | `Choi \| PauliLiouville` | ✓ | PauliLiouville |
139
+ | `PauliLiouville \| StateVector` | ✗ | — |
140
+ | `PauliLiouville \| DensityMatrix` | ✗ | — |
141
+ | `PauliLiouville \| Unitary` | ✓ | PauliLiouville |
142
+ | `PauliLiouville \| SuperOp` | ✓ | SuperOp |
143
+ | `PauliLiouville \| KrausMap` | ✓ | KrausMap |
144
+ | `PauliLiouville \| Choi` | ✓ | Choi |
145
+ | `PauliLiouville \| PauliLiouville` | ✓ | PauliLiouville |
146
+
147
+ ---
148
+
149
+ ### Notes
150
+
151
+ - **Chi** is not included in binary operations because it has no implemented transformations to/from other representations
152
+ - The composition rules follow the principle that when mixing representations, the result uses the "right" operand's representation type
153
+ - State-Operator tensor products (`State | Operator`) return `NotImplemented`
154
+ - Operator-State tensor products (`Operator | State`) return `NotImplemented`
@@ -0,0 +1,65 @@
1
+ # Copyright 2026 Rigetti & Co, LLC.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ [project]
16
+ name = "rigetti-quax"
17
+ version = "0.2.1"
18
+ description = "A high-performance library for quantum information science built on top of JAX"
19
+ authors = [
20
+ {name = "Bram Evert",email = "bevert@rigetti.com"}
21
+ ]
22
+ license = "Apache-2.0"
23
+ readme = "README.md"
24
+ requires-python = ">=3.12"
25
+ dependencies = [
26
+ "jax (>=0.8.2)"
27
+ ]
28
+
29
+ [dependency-groups]
30
+ dev = [
31
+ "pytest (>=9.0.2,<10.0.0)",
32
+ "ruff (>=0.14.10)",
33
+ "qutip (>=5.2.2,<6.0.0)",
34
+ "sphinx (>=8.1.3,<9.0.0)",
35
+ "sphinx-autodoc-typehints (>=2.5.0,<3.0.0)",
36
+ "furo (>=2024.8.6)",
37
+ "sphinx-copybutton (>=0.5.2)",
38
+ "sphinx-design (>=0.6.1)",
39
+ "pyright (>=1.1.407,<2.0.0)",
40
+ "pre-commit (>=4.5.1,<5.0.0)",
41
+ "jupyterlab (>=4.5.2,<5.0.0)",
42
+ "jupyterlab-lsp (>=5.2.0,<6.0.0)",
43
+ "papermill (>=2.6.0)"
44
+ ]
45
+
46
+ [build-system]
47
+ requires = ["poetry-core>=2.0.0,<3.0.0"]
48
+ build-backend = "poetry.core.masonry.api"
49
+
50
+ [tool.poetry]
51
+ packages = [{include = "quax", from = "src"}]
52
+
53
+ [tool.pytest.ini_options]
54
+ testpaths = ["tests"]
55
+ pythonpath = ["src"]
56
+
57
+ [tool.pyright]
58
+ include = ["src", "tests"]
59
+ pythonVersion = "3.12"
60
+ pythonPlatform = "Linux"
61
+ typeCheckingMode = "basic"
62
+
63
+ [tool.ruff]
64
+ line-length = 120
65
+ src = ["src"]
@@ -0,0 +1,249 @@
1
+ # Copyright 2026 Rigetti & Co, LLC.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """JAX-based quantum operator transformations."""
16
+
17
+ from ._apply import (
18
+ apply_choi_to_density_matrix,
19
+ apply_kraus_to_density_matrix,
20
+ apply_kraus_to_state_vector,
21
+ apply_pauli_liouville_to_density_matrix,
22
+ apply_superop_to_density_matrix,
23
+ apply_unitary_to_state_vector,
24
+ compute_choi_observables_from_states,
25
+ compute_kraus_observables_from_states,
26
+ compute_pauli_liouville_observables_from_states,
27
+ compute_superop_observables_from_states,
28
+ partial_trace,
29
+ )
30
+ from ._common_channels import (
31
+ depolarizing_channel_superoperator,
32
+ integrated_thermal_superoperator,
33
+ thermal_relaxation_choi,
34
+ bit_flip_operators,
35
+ phase_flip_operators,
36
+ bitphase_flip_operators,
37
+ dephasing_operators,
38
+ depolarizing_operators,
39
+ relaxation_operators,
40
+ KRAUS_OPS,
41
+ )
42
+ from ._compose import (
43
+ compose_choi,
44
+ compose_kraus,
45
+ compose_pauli_liouville,
46
+ compose_superop,
47
+ compose_unitary,
48
+ )
49
+ from ._distance_metrics import (
50
+ average_fidelity_to_depolarizing_constant,
51
+ average_fidelity_to_process_fidelity,
52
+ depolarizing_constant_to_average_fidelity,
53
+ depolarizing_constant_to_process_fidelity,
54
+ fidelity,
55
+ process_fidelity,
56
+ process_fidelity_to_average_fidelity,
57
+ process_fidelity_to_depolarizing_constant,
58
+ unitarity_to_stochastic_infidelity,
59
+ unitary_entanglement_fidelity,
60
+ )
61
+ from ._power import power_choi, power_kraus, power_pauli_liouville, power_superop, power_unitary
62
+ from ._promotion import promote_state_vector_to_density_matrix
63
+ from ._quantum_objects import (
64
+ Choi,
65
+ DensityMatrix,
66
+ Kraus,
67
+ KrausMap,
68
+ Operator,
69
+ PauliLiouville,
70
+ State,
71
+ StateVector,
72
+ SuperOp,
73
+ SuperOperator,
74
+ Unitary,
75
+ )
76
+ from ._random import (
77
+ ginibre_matrix_complex,
78
+ random_choi_BCSZ,
79
+ random_density_matrix,
80
+ random_state_vector,
81
+ random_unitary,
82
+ )
83
+ from ._state import (
84
+ mixed_state_matrix,
85
+ tensor_density_matrices,
86
+ tensor_state_vectors,
87
+ zero_state_matrix,
88
+ zero_state_vector,
89
+ )
90
+ from ._superoperator_transformations import (
91
+ choi_to_kraus,
92
+ choi_to_pauli_liouville,
93
+ choi_to_superop,
94
+ kraus_to_choi,
95
+ kraus_to_pauli_liouville,
96
+ kraus_to_superop,
97
+ pauli_liouville_to_choi,
98
+ pauli_liouville_to_kraus,
99
+ pauli_liouville_to_superop,
100
+ superop_to_choi,
101
+ superop_to_kraus,
102
+ superop_to_pauli_liouville,
103
+ to_choi,
104
+ to_kraus,
105
+ to_pauli_liouville,
106
+ to_superop,
107
+ unitary_to_choi,
108
+ unitary_to_pauli_liouville,
109
+ unitary_to_superop,
110
+ )
111
+ from ._tensor import (
112
+ tensor_choi,
113
+ tensor_density_matrix,
114
+ tensor_kraus,
115
+ tensor_pauli_liouville,
116
+ tensor_state_vector,
117
+ tensor_superop,
118
+ tensor_unitary,
119
+ )
120
+ from ._validation import (
121
+ is_completely_positive,
122
+ is_cptp,
123
+ is_identity_matrix,
124
+ is_one_design,
125
+ is_trace_preserving,
126
+ is_two_design,
127
+ is_unitary,
128
+ )
129
+
130
+ # Import gates, states and ensembles as submodules
131
+ from . import gates
132
+ from . import states
133
+ from . import ensembles
134
+
135
+ __all__ = [
136
+ # Apply superoperator functions
137
+ "compute_choi_observables_from_states",
138
+ "compute_kraus_observables_from_states",
139
+ "compute_pauli_liouville_observables_from_states",
140
+ "compute_superop_observables_from_states",
141
+ "partial_trace",
142
+ "apply_choi_to_density_matrix",
143
+ "apply_kraus_to_density_matrix",
144
+ "apply_pauli_liouville_to_density_matrix",
145
+ "apply_superop_to_density_matrix",
146
+ "apply_unitary_to_state_vector",
147
+ "apply_kraus_to_state_vector",
148
+ # Common channels
149
+ "depolarizing_channel_superoperator",
150
+ "integrated_thermal_superoperator",
151
+ "thermal_relaxation_choi",
152
+ "bit_flip_operators",
153
+ "phase_flip_operators",
154
+ "bitphase_flip_operators",
155
+ "dephasing_operators",
156
+ "depolarizing_operators",
157
+ "relaxation_operators",
158
+ "KRAUS_OPS",
159
+ # Compose quantum objects
160
+ "compose_kraus",
161
+ "compose_choi",
162
+ "compose_pauli_liouville",
163
+ "compose_superop",
164
+ "compose_unitary",
165
+ # Tensor quantum objects
166
+ "tensor_kraus",
167
+ "tensor_choi",
168
+ "tensor_pauli_liouville",
169
+ "tensor_superop",
170
+ "tensor_unitary",
171
+ "tensor_state_vector",
172
+ "tensor_density_matrix",
173
+ # Superoperator transformations
174
+ "choi_to_kraus",
175
+ "choi_to_pauli_liouville",
176
+ "choi_to_superop",
177
+ "kraus_to_choi",
178
+ "kraus_to_pauli_liouville",
179
+ "kraus_to_superop",
180
+ "pauli_liouville_to_choi",
181
+ "pauli_liouville_to_kraus",
182
+ "pauli_liouville_to_superop",
183
+ "superop_to_choi",
184
+ "superop_to_kraus",
185
+ "superop_to_pauli_liouville",
186
+ "to_choi",
187
+ "to_pauli_liouville",
188
+ "to_superop",
189
+ "to_kraus",
190
+ "unitary_to_choi",
191
+ "unitary_to_pauli_liouville",
192
+ "unitary_to_superop",
193
+ # Distance metrics
194
+ "fidelity",
195
+ "process_fidelity",
196
+ "to_choi",
197
+ "unitary_entanglement_fidelity",
198
+ "depolarizing_constant_to_average_fidelity",
199
+ "average_fidelity_to_depolarizing_constant",
200
+ "depolarizing_constant_to_process_fidelity",
201
+ "process_fidelity_to_depolarizing_constant",
202
+ "average_fidelity_to_process_fidelity",
203
+ "process_fidelity_to_average_fidelity",
204
+ "unitarity_to_stochastic_infidelity",
205
+ # Types
206
+ "Choi",
207
+ "Kraus",
208
+ "SuperOp",
209
+ "KrausMap",
210
+ "PauliLiouville",
211
+ "SuperOperator",
212
+ "Unitary",
213
+ "DensityMatrix",
214
+ "StateVector",
215
+ "State",
216
+ "Operator",
217
+ # State functions
218
+ "zero_state_matrix",
219
+ "zero_state_vector",
220
+ "mixed_state_matrix",
221
+ "tensor_density_matrices",
222
+ "tensor_state_vectors",
223
+ # Random functions
224
+ "ginibre_matrix_complex",
225
+ "random_density_matrix",
226
+ "random_unitary",
227
+ "random_choi_BCSZ",
228
+ "random_state_vector",
229
+ # Validation functions
230
+ "is_unitary",
231
+ "is_one_design",
232
+ "is_two_design",
233
+ "is_identity_matrix",
234
+ "is_cptp",
235
+ "is_completely_positive",
236
+ "is_trace_preserving",
237
+ # Promotion functions
238
+ "promote_state_vector_to_density_matrix",
239
+ # Power functions
240
+ "power_choi",
241
+ "power_kraus",
242
+ "power_pauli_liouville",
243
+ "power_superop",
244
+ "power_unitary",
245
+ # Submodules
246
+ "gates",
247
+ "states",
248
+ "ensembles",
249
+ ]