1"""
2A collection of methods for packaging and unpackaging varaibles into a single
3tensor.
4"""
5
6import typing
7
8from typing import Iterable
9
10import numpy as np
11
[docs]
12def package_complex(x: np.ndarray[float]) -> typing.Any:
13 """Takes the last axis of a real tensor and takes the first entry as the
14 real part and the second as the imaginary.
15
16 Parameters
17 ----------
18 x: NDArray[Shape[s := Any_Shape, 2], float]
19 The tensor to reshape.
20
21 Returns
22 -------
23 NDArray[Shape[s], complex]
24 The reshaped tensor.
25
26 Methods
27 -------
28 package_complex.inverse
29 Takes a complex tensor and splits it into the real and imaginary parts
30 as the final axis.
31
32 PARAMETERS:
33 * **z** (*NDArray[Shape[s := Any_Shape], complex]*) —
34 The tensor to reshape.
35
36 RETURNS:
37 The reshaped tensor.
38
39 RETURN TYPE:
40 NDArray[Shape[s, 2], float]
41 package_complex.compose
42 Composes :func:`package_complex()` with another
43 :class:`.InvertibleFunction` to create a new
44 :class:`.InvertibleFunction` along with the composed inverse. That is
45 the following assertions should hold::
46
47 assert package_complex.compose(g)(x, *g_args, **g_kwargs) \
48 == package_complex(g(x, *g_args, **g_kwargs))
49
50 for all inputs ``x``.
51
52 PARAMETERS:
53 * **inner_invertible_function** (*InvertibleFunction*) —
54 The :class:`.InvertibleFunction` be be called first. The output of
55 this :class:`.InvertibleFunction` will be passed to
56 :func:`package_complex()`.
57
58 RETURNS:
59 A new :class:`.InvertibleFunction` that is the composition of the
60 two functions.
61
62 RETURN TYPE:
63 :class:`.InvertibleFunction`
64 """
65 ...
[docs]
66def unpackage_complex(z: np.ndarray[complex]) -> typing.Any:
67 """Takes a complex tensor and splits it into the real and imaginary parts as
68 the final axis.
69
70 Parameters
71 ----------
72 z: NDArray[Shape[s := Any_Shape], complex]
73 The tensor to reshape.
74
75 Returns
76 -------
77 NDArray[Shape[s, 2], float]
78 The reshaped tensor.
79
80 Methods
81 -------
82 unpackage_complex.inverse
83 Takes the last axis of a real tensor and takes the first entry as the
84 real part and the second as the imaginary.
85
86 PARAMETERS:
87 * **x** (*NDArray[Shape[s := Any_Shape, 2], float]*) —
88 The tensor to reshape.
89
90 RETURNS:
91 The reshaped tensor.
92
93 RETURN TYPE:
94 NDArray[Shape[s], complex]
95 unpackage_complex.compose
96 Composes the :func:`unpackage_complex()` with another
97 :class:`.InvertibleFunction` to create a new
98 :class:`.InvertibleFunction` along with the composed inverse.
99 That is the following assertions should hold::
100
101 assert package_complex.compose(g)(z, *g_args, **g_kwargs) \
102 == package_complex(g(z, *g_args, **g_kwargs))
103
104 for all inputs ``z``.
105
106 PARAMETERS:
107 * **inner_invertible_function** (*InvertibleFunction*) —
108 The :class:`.InvertibleFunction` be be called first. The output of
109 this :class:`.InvertibleFunction` will be passed to
110 :func:`unpackage_complex()`.
111
112 RETURNS:
113 A new :class:`.InvertibleFunction` that is the composition of the
114 two functions.
115
116 RETURN TYPE:
117 :class:`.InvertibleFunction`
118 """
119 ...
120
[docs]
121class ViewPop():
122 """
123 A class that allows for a view of an iterable to be popped.
124
125 Example
126 -------
127 >>> x = [1, 2, 3, 4, 5, 6]
128 >>> v = ViewPop(x)
129 >>> v(2)
130 [1, 2]
131 >>> v(2)
132 [3, 4]
133 >>> v(1)
134 [5]
135 >>> v(2)
136 [6]
137 """
138
139 iterator: Iterable
140 """
141 The iterator to be viewed
142 """
143
144 index: int
145 """
146 The current index of the iterator
147 """
148
[docs]
149 def __init__(self, iterator: Iterable):
150 """
151 Initialises the view of the `iterator`.
152
153 Parameters
154 ----------
155 iterator : Iterable
156 The iterator to be viewed
157 """
158 ...
159 def __call__(self, number_of_elements_to_pop: int = 1) -> Iterable:
160 """
161 Pops a view of the `iterator` of the given size.
162
163 Parameters
164 ----------
165 number_of_elements_to_pop : int
166 The number of elements to pop from the iterator
167
168 Returns
169 -------
170 Iterable
171 A view of the iterator of the given size
172 """
173 ...
174
[docs]
175def unpack(x: Iterable, shapes: Iterable[Iterable[int]]) -> list:
176 """
177 Unpacks an input iterable `x` into
178 `TensorFlow <https://www.tensorflow.org>`__ tensors of the given shapes.
179
180 Parameters
181 ----------
182 x: NDArray[Shape[length], complex]
183 The iterable to be unpacked.
184 shapes: Iterable[Iterable[int]]
185 The shapes of the tensors to be unpacked into.
186
187 Returns
188 -------
189 list
190 A list of `TensorFlow <https://www.tensorflow.org>`__ tensors of the
191 given shapes.
192
193 Methods
194 -------
195 unpack.inverse
196 Packs the list of input tensors into a single
197 `TensorFlow <https://www.tensorflow.org>`__ tensor.
198
199 PARAMETERS:
200 * **tensors** (*list*) —
201 The `TensorFlow <https://www.tensorflow.org>`__ tensors to be
202 packed
203
204 RETURNS:
205 A `TensorFlow <https://www.tensorflow.org>`__ tensor with 1 axis
206 consisting of all the input tensors flattened and concatenated.
207
208 RETURN TYPE:
209 Any
210 unpack.specify_parameters
211 Allows the ``shapes`` to be pre-specified. This removes shapes from the
212 call signature.
213
214 PARAMETERS:
215 * **shapes** (*Iterable[Iterable[int]], optional*) —
216 The shapes of the tensors to be unpacked into.
217
218 RETURNS:
219 A new :class:`.InvertibleFunction` with the specified parameters
220 pre-specified.
221
222 RETURN TYPE:
223 InvertibleFunction
224 unpack.compose
225 Composes the :func:`unpack()` with another :class:`.InvertibleFunction`
226 to create a new :class:`.InvertibleFunction` along with the composed
227 inverse. That is the following assertions should hold::
228
229 assert unpack.compose(g, *args, **kwargs)(x, *g_args, **g_kwargs) \
230 == unpack(g(x, *g_args, **g_kwargs), *args, **kwargs)
231
232 for all inputs ``x``.
233
234 PARAMETERS:
235 * **inner_invertible_function** (*InvertibleFunction*) —
236 The :class:`.InvertibleFunction` be be called first. The output of
237 this :class:`.InvertibleFunction` will be passed to
238 :func:`unpack()`.
239 * **shapes** (*Iterable[Iterable[int]]*) —
240 The shapes of the tensors to be unpacked into.
241
242 RETURNS:
243 A new :class:`.InvertibleFunction` that is the composition of the
244 two functions.
245
246 RETURN TYPE:
247 InvertibleFunction
248 """
249 ...
[docs]
250def pack(tensors: list) -> typing.Any:
251 """
252 Packs the list of input tensors into a single
253 `TensorFlow <https://www.tensorflow.org>`__ tensor.
254
255 Parameters
256 ----------
257 tensors: list
258 The list of `TensorFlow <https://www.tensorflow.org>`__ tensors to be
259 packed
260
261 Returns
262 -------
263 `TensorFlow <https://www.tensorflow.org>`__ Tensor
264 A `TensorFlow <https://www.tensorflow.org>`__ tensor with 1 axis
265 consisting of all the input tensors flattened and concatenated.
266
267 Methods
268 -------
269 pack.inverse
270 Unpacks an input iterable `x` into
271 `TensorFlow <https://www.tensorflow.org>`__ tensors of the given shapes.
272
273 PARAMETERS:
274 * **x** (*Iterable*) —
275 The iterable to be unpacked.
276 * **shapes** (*Iterable[Iterable[int]]*) —
277 The shapes of the tensors to be unpacked into.
278
279 RETURNS:
280 A `TensorFlow <https://www.tensorflow.org>`__ tensor with 1 axis
281 consisting of all the input tensors flattened and concatenated.
282
283 RETURN TYPE:
284 list
285 pack.specify_parameters
286 Allows the ``shapes`` to be pre-specified. This removes shapes from the
287 call signature of :func:`inverse()`.
288
289 PARAMETERS:
290 * **shapes** (*Iterable[Iterable[int]], optional*) —
291 The shapes of the tensors to be unpacked into.
292
293 RETURNS:
294 A new :class:`.InvertibleFunction` with the specified parameters
295 pre-specified.
296
297 RETURN TYPE:
298 InvertibleFunction
299 pack.compose
300 Composes the :func:`pack()` with another :class:`.InvertibleFunction`
301 to create a new :class:`.InvertibleFunction` along with the composed
302 inverse. That is the following assertions should hold::
303
304 assert pack.compose(g, *args, **kwargs)(x, *g_args, **g_kwargs) \
305 == pack(g(x, *g_args, **g_kwargs), *args, **kwargs)
306
307 for all inputs ``x``.
308
309 PARAMETERS:
310 * **inner_invertible_function** (*InvertibleFunction*) —
311 The :class:`.InvertibleFunction` be be called first. The output of
312 this :class:`.InvertibleFunction` will be passed to
313 :func:`pack()`.
314 * **shapes** (*Iterable[Iterable[int]]*) —
315 The shapes of the tensors to be unpacked into.
316
317 RETURNS:
318 A new :class:`.InvertibleFunction` that is the composition of the
319 two functions.
320
321 RETURN TYPE:
322 InvertibleFunction
323 """
324 ...