qudit_expr/shape.rs
1use serde::{Deserialize, Serialize};
2
3use crate::index::{IndexDirection, TensorIndex};
4use std::ops::Add;
5
6/// Represents the shape of a tensor as it will be generated.
7///
8/// While tensors can conceptually have rank larger than four, even infinite,
9/// tensors in the OpenQudit Expression library are generated into a buffer
10/// indexed by 0, 1, 2, 3, or 4 physical dimensions.
11#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
12pub enum GenerationShape {
13 /// A 0-dimensional tensor (a single value).
14 Scalar,
15
16 /// A 1-dimensional tensor with `nelems` elements.
17 Vector(usize),
18
19 /// A 2-dimensional tensor (matrix) with `nrows` rows and `ncols` columns.
20 Matrix(usize, usize),
21
22 /// A 3-dimensional tensor with `nmats` matrices, each of `nrows` rows and `ncols` columns.
23 Tensor3D(usize, usize, usize),
24
25 /// A 4-dimensional tensor usually for derivatives (ntens, nmats, nrows, ncols)
26 Tensor4D(usize, usize, usize, usize),
27}
28
29impl std::fmt::Debug for GenerationShape {
30 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31 match self {
32 GenerationShape::Scalar => write!(f, "Scalar"),
33 GenerationShape::Vector(nelems) => write!(f, "Vector({})", nelems),
34 GenerationShape::Matrix(nrows, ncols) => write!(f, "Matrix({}, {})", nrows, ncols),
35 GenerationShape::Tensor3D(nmats, nrows, ncols) => {
36 write!(f, "Tensor3D({}, {}, {})", nmats, nrows, ncols)
37 }
38 GenerationShape::Tensor4D(ntens, nmats, nrows, ncols) => {
39 write!(f, "Tensor4D({}, {}, {}, {})", ntens, nmats, nrows, ncols)
40 }
41 }
42 }
43}
44
45impl GenerationShape {
46 /// Calculates the total number of elements in a tensor with this shape.
47 ///
48 /// # Returns
49 /// The total number of elements as `usize`.
50 pub fn num_elements(&self) -> usize {
51 match self {
52 GenerationShape::Scalar => 1,
53 GenerationShape::Vector(nelems) => *nelems,
54 GenerationShape::Matrix(nrows, ncols) => nrows * ncols,
55 GenerationShape::Tensor3D(nmats, nrows, ncols) => nmats * nrows * ncols,
56 GenerationShape::Tensor4D(ntens, nmats, nrows, ncols) => ntens * nmats * nrows * ncols,
57 }
58 }
59
60 /// Determines the shape of the derivative of a tensor with respect to `num_params` parameters.
61 ///
62 /// This method effectively prepends `num_params` to the current tensor's dimensions.
63 /// For example, the derivative of a `Scalar` with respect to `num_params` becomes a `Vector(num_params)`.
64 /// The derivative of a `Matrix(R, C)` becomes a `Tensor3D(num_params, R, C)`.
65 ///
66 /// # Arguments
67 /// * `num_params` - The number of parameters in the gradient.
68 ///
69 /// # Returns
70 /// A new `GenerationShape` representing the shape of the gradient.
71 ///
72 /// # See Also
73 /// - `[hessian_shape]` For the shape of a hessian tensor.
74 pub fn gradient_shape(&self, num_params: usize) -> Self {
75 match self {
76 GenerationShape::Scalar => GenerationShape::Vector(num_params),
77 GenerationShape::Vector(nelems) => GenerationShape::Matrix(num_params, *nelems),
78 GenerationShape::Matrix(nrows, ncols) => {
79 GenerationShape::Tensor3D(num_params, *nrows, *ncols)
80 }
81 GenerationShape::Tensor3D(nmats, nrows, ncols) => {
82 GenerationShape::Tensor4D(num_params, *nmats, *nrows, *ncols)
83 }
84 GenerationShape::Tensor4D(_, _, _, _) => {
85 panic!("Unable to find shape for gradient of 4D tensor.")
86 }
87 }
88 }
89
90 /// Determine the hessian shape of a tensor with this shape that has `num_params` parameters.
91 ///
92 /// # Arguments
93 /// * `num_params` - The number of parameters in the hessian.
94 ///
95 /// # Returns
96 /// A new `GenerationShape` representing the shape of the hessian.
97 ///
98 /// # See Also
99 /// - `[gradient_shape]` For the shape of a gradient tensor.
100 pub fn hessian_shape(&self, num_params: usize) -> Self {
101 let sym_sq_size = num_params * (num_params - 1) / 2; // TODO: Is it +1 or -1
102 match self {
103 GenerationShape::Scalar => GenerationShape::Vector(sym_sq_size),
104 GenerationShape::Vector(nelems) => GenerationShape::Matrix(sym_sq_size, *nelems),
105 GenerationShape::Matrix(nrows, ncols) => {
106 GenerationShape::Tensor3D(sym_sq_size, *nrows, *ncols)
107 }
108 GenerationShape::Tensor3D(nmats, nrows, ncols) => {
109 GenerationShape::Tensor4D(sym_sq_size, *nmats, *nrows, *ncols)
110 }
111 GenerationShape::Tensor4D(_, _, _, _) => {
112 panic!("Unable to find shape for Hessian of 4D tensor.")
113 }
114 }
115 }
116
117 /// Converts the tensor shape object to a vector of integers.
118 ///
119 /// # Returns
120 ///
121 /// A `Vec<usize>` containing the dimensions of the shape.
122 ///
123 /// # Examples
124 /// ```
125 /// use qudit_expr::GenerationShape;
126 ///
127 /// let scalar_shape = GenerationShape::Scalar;
128 /// assert_eq!(scalar_shape.to_vec(), Vec::<usize>::new());
129 ///
130 /// let vector_shape = GenerationShape::Vector(5);
131 /// assert_eq!(vector_shape.to_vec(), vec![5]);
132 ///
133 /// let matrix_shape = GenerationShape::Matrix(2, 3);
134 /// assert_eq!(matrix_shape.to_vec(), vec![2, 3]);
135 /// ```
136 pub fn to_vec(&self) -> Vec<usize> {
137 match self {
138 GenerationShape::Scalar => vec![],
139 GenerationShape::Vector(nelems) => vec![*nelems],
140 GenerationShape::Matrix(nrows, ncols) => vec![*nrows, *ncols],
141 GenerationShape::Tensor3D(nmats, nrows, ncols) => vec![*nmats, *nrows, *ncols],
142 GenerationShape::Tensor4D(ntens, nmats, nrows, ncols) => {
143 vec![*ntens, *nmats, *nrows, *ncols]
144 }
145 }
146 }
147
148 /// Checks if the current `GenerationShape` is strictly a scalar variant.
149 pub fn is_scalar(&self) -> bool {
150 matches!(self, GenerationShape::Scalar)
151 }
152
153 /// Checks if the current `GenerationShape` is strictly a vector variant.
154 pub fn is_vector(&self) -> bool {
155 matches!(self, GenerationShape::Vector(_))
156 }
157
158 /// Checks if the current `GenerationShape` is strictly a matrix variant.
159 pub fn is_matrix(&self) -> bool {
160 matches!(self, GenerationShape::Matrix(_, _))
161 }
162
163 /// Checks if the current `GenerationShape` is strictly a tensor3D variant.
164 pub fn is_tensor3d(&self) -> bool {
165 matches!(self, GenerationShape::Tensor3D(_, _, _))
166 }
167
168 /// Checks if the current `GenerationShape` is strictly a tensor4D variant.
169 pub fn is_tensor4d(&self) -> bool {
170 matches!(self, GenerationShape::Tensor4D(_, _, _, _))
171 }
172
173 /// Check if there is only one element.
174 ///
175 /// # Returns
176 /// `true` if the shape can be treated as a scalar, `false` otherwise.
177 ///
178 /// # Examples
179 /// ```
180 /// use qudit_expr::GenerationShape;
181 ///
182 /// let test_scalar = GenerationShape::Scalar;
183 /// let test_vector = GenerationShape::Vector(1);
184 /// let test_matrix = GenerationShape::Matrix(1, 1);
185 /// let test_tensor3d = GenerationShape::Tensor3D(1, 1, 1);
186 ///
187 /// let test_vector_2 = GenerationShape::Vector(9);
188 /// let test_matrix_2 = GenerationShape::Matrix(9, 9);
189 /// let test_tensor3d_2 = GenerationShape::Tensor3D(1, 9, 9);
190 ///
191 /// assert!(test_scalar.is_0d());
192 /// assert!(test_vector.is_0d());
193 /// assert!(test_matrix.is_0d());
194 /// assert!(test_tensor3d.is_0d());
195 ///
196 /// assert_eq!(test_vector_2.is_0d(), false);
197 /// assert_eq!(test_matrix_2.is_0d(), false);
198 /// assert_eq!(test_tensor3d_2.is_0d(), false);
199 /// ```
200 pub fn is_0d(&self) -> bool {
201 self.num_elements() == 1
202 }
203
204 /// Check if the shape can be conceptually treated as a 1d tensor.
205 ///
206 /// # Returns
207 /// `true` if the shape has exactly one dimension with 1 or more elements.
208 ///
209 /// # Examples
210 /// ```
211 /// use qudit_expr::GenerationShape;
212 ///
213 /// let test_vector = GenerationShape::Vector(9);
214 /// let test_matrix = GenerationShape::Matrix(1, 9);
215 /// let test_tensor3d = GenerationShape::Tensor3D(1, 1, 9);
216 ///
217 /// let test_scalar = GenerationShape::Scalar;
218 /// let test_matrix_2 = GenerationShape::Matrix(9, 9);
219 /// let test_tensor3d_2 = GenerationShape::Tensor3D(1, 9, 9);
220 ///
221 /// assert!(test_vector.is_1d());
222 /// assert!(test_matrix.is_1d());
223 /// assert!(test_tensor3d.is_1d());
224 ///
225 /// assert_eq!(test_scalar.is_1d(), false);
226 /// assert_eq!(test_matrix_2.is_1d(), false);
227 /// assert_eq!(test_tensor3d_2.is_1d(), false);
228 /// ```
229 pub fn is_1d(&self) -> bool {
230 match self {
231 GenerationShape::Scalar => false,
232 GenerationShape::Vector(_) => true,
233 GenerationShape::Matrix(nrows, ncols) => *nrows == 1 || *ncols == 1,
234 GenerationShape::Tensor3D(nmats, nrows, ncols) => {
235 let non_one_count = [*nmats, *nrows, *ncols].iter().filter(|&&d| d > 1).count();
236 non_one_count == 1
237 }
238 GenerationShape::Tensor4D(ntens, nmats, nrows, ncols) => {
239 let non_one_count = [*ntens, *nmats, *nrows, *ncols]
240 .iter()
241 .filter(|&&d| d > 1)
242 .count();
243 non_one_count == 1
244 }
245 }
246 }
247
248 /// Checks if the current `GenerationShape` can be conceptually treated as a 2-dimensional matrix.
249 /// This is true for `GenerationShape` variants with a dimensionality of at least 2, with
250 /// any additional dimensions having size 1.
251 ///
252 /// # Returns
253 /// `true` if the shape can be treated as a matrix, `false` otherwise.
254 ///
255 /// # Examples
256 /// ```
257 /// use qudit_expr::GenerationShape;
258 ///
259 /// let test_scalar = GenerationShape::Scalar;
260 /// let test_vector = GenerationShape::Vector(1);
261 /// let test_tensor3d_2 = GenerationShape::Tensor3D(9, 9, 9);
262 /// let test_tensor_nd_2 = GenerationShape::Tensor4D(1, 9, 9, 9);
263 ///
264 /// let test_matrix = GenerationShape::Matrix(9, 9);
265 /// let test_tensor3d = GenerationShape::Tensor3D(1, 9, 9);
266 /// let test_tensor_nd = GenerationShape::Tensor4D(1, 1, 9, 9);
267 ///
268 /// assert_eq!(test_scalar.is_2d(), false);
269 /// assert_eq!(test_vector.is_2d(), false);
270 /// assert_eq!(test_tensor3d_2.is_2d(), false);
271 /// assert_eq!(test_tensor_nd_2.is_2d(), false);
272 ///
273 /// assert_eq!(test_matrix.is_2d(), true);
274 /// assert_eq!(test_tensor3d.is_2d(), true);
275 /// assert_eq!(test_tensor_nd.is_2d(), true);
276 /// ```
277 pub fn is_2d(&self) -> bool {
278 match self {
279 GenerationShape::Scalar => false,
280 GenerationShape::Vector(_) => false,
281 GenerationShape::Matrix(_, _) => true,
282 // A Tensor3D can be seen as a matrix if it's essentially a stack of column vectors,
283 // or perhaps if it represents a single matrix (nmats=1).
284 // The current implementation checks if ncols is 1, implying it's a stack of column vectors.
285 GenerationShape::Tensor3D(nmats, _, _) => *nmats == 1,
286 GenerationShape::Tensor4D(ntens, nmats, _, _) => *ntens == 1 && *nmats == 1,
287 }
288 }
289
290 /// Checks if the current `GenerationShape` can be conceptually treated as a 3D tensor.
291 /// This is true for `GenerationShape` variants with a dimensionality of at least 3, with
292 /// any additional dimensions having size 1.
293 ///
294 /// # Returns
295 /// `true` if the shape can be treated as a 3D tensor, `false` otherwise.
296 ///
297 /// # Examples
298 /// ```
299 /// use qudit_expr::GenerationShape;
300 ///
301 /// let test_scalar = GenerationShape::Scalar;
302 /// let test_vector = GenerationShape::Vector(1);
303 /// let test_matrix = GenerationShape::Matrix(1, 1);
304 /// let test_tensor_nd_2 = GenerationShape::Tensor4D(9, 1, 9, 9);
305 ///
306 /// let test_tensor3d = GenerationShape::Tensor3D(9, 9, 9);
307 /// let test_tensor_nd = GenerationShape::Tensor4D(1, 9, 9, 9);
308 ///
309 /// assert_eq!(test_scalar.is_3d(), false);
310 /// assert_eq!(test_vector.is_3d(), false);
311 /// assert_eq!(test_matrix.is_3d(), false);
312 /// assert_eq!(test_tensor_nd_2.is_3d(), false);
313 ///
314 /// assert_eq!(test_tensor3d.is_3d(), true);
315 /// assert_eq!(test_tensor_nd.is_3d(), true);
316 /// ```
317 pub fn is_3d(&self) -> bool {
318 match self {
319 GenerationShape::Scalar => false,
320 GenerationShape::Vector(_) => false,
321 GenerationShape::Matrix(_, _) => false,
322 GenerationShape::Tensor3D(_, _, _) => true,
323 GenerationShape::Tensor4D(ntens, _, _, _) => *ntens == 1,
324 }
325 }
326
327 /// Checks if the current `GenerationShape` can be conceptually treated as a 4D tensor.
328 /// This is true for `GenerationShape` variants with a dimensionality of at least 4, with
329 /// any additional dimensions having size 1.
330 ///
331 /// # Returns
332 /// `true` if the shape can be treated as a 4D tensor, `false` otherwise.
333 ///
334 /// # Examples
335 /// ```
336 /// use qudit_expr::GenerationShape;
337 ///
338 /// let test_scalar = GenerationShape::Scalar;
339 /// let test_vector = GenerationShape::Vector(1);
340 /// let test_matrix = GenerationShape::Matrix(1, 1);
341 /// let test_tensor3d = GenerationShape::Tensor3D(1, 1, 1);
342 ///
343 /// let test_tensor4d = GenerationShape::Tensor4D(9, 9, 9, 9);
344 /// let test_tensor_nd = GenerationShape::Tensor4D(1, 9, 9, 9);
345 ///
346 /// assert_eq!(test_scalar.is_4d(), false);
347 /// assert_eq!(test_vector.is_4d(), false);
348 /// assert_eq!(test_matrix.is_4d(), false);
349 /// assert_eq!(test_tensor3d.is_4d(), false);
350 ///
351 /// assert_eq!(test_tensor4d.is_4d(), true);
352 /// assert_eq!(test_tensor_nd.is_4d(), true);
353 /// ```
354 pub fn is_4d(&self) -> bool {
355 match self {
356 GenerationShape::Scalar => false,
357 GenerationShape::Vector(_) => false,
358 GenerationShape::Matrix(_, _) => false,
359 GenerationShape::Tensor3D(_, _, _) => false,
360 GenerationShape::Tensor4D(_, _, _, _) => true,
361 }
362 }
363
364 /// Returns the number of columns for the current shape.
365 ///
366 /// # Returns
367 /// The number of columns.
368 ///
369 /// # Examples
370 /// ```
371 /// use qudit_expr::GenerationShape;
372 /// let matrix_shape = GenerationShape::Matrix(2, 3);
373 /// assert_eq!(matrix_shape.ncols(), 3);
374 /// ```
375 pub fn ncols(&self) -> usize {
376 match self {
377 GenerationShape::Scalar => 1,
378 GenerationShape::Vector(ncols) => *ncols,
379 GenerationShape::Matrix(_, ncols) => *ncols,
380 GenerationShape::Tensor3D(_, _, ncols) => *ncols,
381 GenerationShape::Tensor4D(_, _, _, ncols) => *ncols,
382 }
383 }
384
385 /// Returns the number of rows for the current shape.
386 ///
387 /// # Returns
388 /// The number of rows.
389 ///
390 /// # Examples
391 /// ```
392 /// use qudit_expr::GenerationShape;
393 /// let matrix_shape = GenerationShape::Matrix(2, 3);
394 /// assert_eq!(matrix_shape.nrows(), 2);
395 /// ```
396 pub fn nrows(&self) -> usize {
397 match self {
398 GenerationShape::Scalar => 1,
399 GenerationShape::Vector(_) => 1,
400 GenerationShape::Matrix(nrows, _) => *nrows,
401 GenerationShape::Tensor3D(_, nrows, _) => *nrows,
402 GenerationShape::Tensor4D(_, _, nrows, _) => *nrows,
403 }
404 }
405
406 /// Returns the number of matrices for the current shape.
407 ///
408 /// # Returns
409 /// The number of matrices.
410 ///
411 /// # Examples
412 /// ```
413 /// use qudit_expr::GenerationShape;
414 /// let tensor3d_shape = GenerationShape::Tensor3D(5, 2, 3);
415 /// assert_eq!(tensor3d_shape.nmats(), 5);
416 /// ```
417 pub fn nmats(&self) -> usize {
418 match self {
419 GenerationShape::Scalar => 1,
420 GenerationShape::Vector(_) => 1,
421 GenerationShape::Matrix(_, _) => 1,
422 GenerationShape::Tensor3D(nmats, _, _) => *nmats,
423 GenerationShape::Tensor4D(_, nmats, _, _) => *nmats,
424 }
425 }
426
427 /// Returns the number of tensors (in the first dimension) for the current shape.
428 ///
429 /// # Returns
430 /// The number of tensors.
431 ///
432 /// # Examples
433 /// ```
434 /// use qudit_expr::GenerationShape;
435 /// let tensor4d_shape = GenerationShape::Tensor4D(7, 5, 2, 3);
436 /// assert_eq!(tensor4d_shape.ntens(), 7);
437 /// ```
438 pub fn ntens(&self) -> usize {
439 match self {
440 GenerationShape::Scalar => 1,
441 GenerationShape::Vector(_) => 1,
442 GenerationShape::Matrix(_, _) => 1,
443 GenerationShape::Tensor3D(_, _, _) => 1,
444 GenerationShape::Tensor4D(ntens, _, _, _) => *ntens,
445 }
446 }
447
448 pub fn calculate_directions(&self, index_sizes: &[usize]) -> Vec<IndexDirection> {
449 match self {
450 GenerationShape::Scalar => vec![],
451 GenerationShape::Vector(_) => vec![IndexDirection::Input; index_sizes.len()],
452 GenerationShape::Matrix(nrows, _) => {
453 let mut index_size_acm = 1usize;
454 let mut index_iter = 0;
455 let mut index_directions = vec![];
456 while index_iter < index_sizes.len() && index_size_acm < *nrows {
457 index_size_acm *= index_sizes[index_iter];
458 index_directions.push(IndexDirection::Output);
459 index_iter += 1;
460 }
461 while index_iter < index_sizes.len() {
462 index_directions.push(IndexDirection::Input);
463 index_iter += 1;
464 }
465 index_directions
466 }
467 GenerationShape::Tensor3D(nmats, nrows, _) => {
468 let mut index_size_acm = 1usize;
469 let mut index_iter = 0;
470 let mut index_directions = vec![];
471 while index_iter < index_sizes.len() && index_size_acm < *nmats {
472 index_size_acm *= index_sizes[index_iter];
473 index_directions.push(IndexDirection::Batch);
474 index_iter += 1;
475 }
476 index_size_acm = 1usize; // Reset to calculate for nrows
477 while index_iter < index_sizes.len() && index_size_acm < *nrows {
478 index_size_acm *= index_sizes[index_iter];
479 index_directions.push(IndexDirection::Output);
480 index_iter += 1;
481 }
482 while index_iter < index_sizes.len() {
483 index_directions.push(IndexDirection::Input);
484 index_iter += 1;
485 }
486 index_directions
487 }
488 GenerationShape::Tensor4D(_, _, _, _) => {
489 todo!()
490 }
491 }
492 }
493}
494
495// impl From<Vec<TensorIndex>> for GenerationShape {
496// fn from(indices: Vec<TensorIndex>) -> Self {
497// GenerationShape::from(indices.as_slice())
498// }
499// }
500
501impl<I: AsRef<[TensorIndex]>> From<I> for GenerationShape {
502 fn from(indices: I) -> Self {
503 let indices = indices.as_ref();
504 let mut dimensions = [1, 1, 1, 1];
505 for index in indices.iter() {
506 match index.direction() {
507 IndexDirection::Derivative => dimensions[0] *= index.index_size(),
508 IndexDirection::Batch => dimensions[1] *= index.index_size(),
509 IndexDirection::Output => dimensions[2] *= index.index_size(),
510 IndexDirection::Input => dimensions[3] *= index.index_size(),
511 }
512 }
513
514 match dimensions {
515 [1, 1, 1, 1] => GenerationShape::Scalar,
516 [1, 1, 1, nelems] => GenerationShape::Vector(nelems),
517 [1, 1, nrows, ncols] => GenerationShape::Matrix(nrows, ncols),
518 [1, nmats, nrows, ncols] => GenerationShape::Tensor3D(nmats, nrows, ncols),
519 [ntens, nmats, nrows, ncols] => GenerationShape::Tensor4D(ntens, nmats, nrows, ncols),
520 }
521 }
522}
523
524impl Add for GenerationShape {
525 type Output = Self;
526
527 fn add(self, other: Self) -> Self::Output {
528 match (self, other) {
529 // TODO: re-evaluate this...
530 (GenerationShape::Scalar, other_shape) => other_shape,
531 (other_shape, GenerationShape::Scalar) => other_shape,
532 (GenerationShape::Vector(s_nelems), GenerationShape::Vector(o_nelems)) => {
533 GenerationShape::Vector(s_nelems + o_nelems)
534 }
535 (
536 GenerationShape::Matrix(s_nrows, s_ncols),
537 GenerationShape::Matrix(o_nrows, o_ncols),
538 ) => GenerationShape::Matrix(s_nrows + o_nrows, s_ncols + o_ncols),
539 (
540 GenerationShape::Tensor3D(s_nmats, s_nrows, s_ncols),
541 GenerationShape::Tensor3D(o_nmats, o_nrows, o_ncols),
542 ) => GenerationShape::Tensor3D(s_nmats + o_nmats, s_nrows + o_nrows, s_ncols + o_ncols),
543 (
544 GenerationShape::Tensor4D(s_ntens, s_nmats, s_nrows, s_ncols),
545 GenerationShape::Tensor4D(o_ntens, o_nmats, o_nrows, o_ncols),
546 ) => GenerationShape::Tensor4D(
547 s_ntens + o_ntens,
548 s_nmats + o_nmats,
549 s_nrows + o_nrows,
550 s_ncols + o_ncols,
551 ),
552 _ => {
553 panic!("Cannot add tensors of different fundamental shapes or incompatible ranks.")
554 }
555 }
556 }
557}