qudit-tensor 0.3.3

Accelerated and Extensible Quantum Library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use qudit_core::ParamInfo;
use qudit_expr::ExpressionId;
use qudit_expr::GenerationShape;

use qudit_expr::index::IndexDirection;
use qudit_expr::index::TensorIndex;

#[derive(Debug, Clone)]
pub struct QuditTensor {
    pub indices: Vec<TensorIndex>,
    pub expression: ExpressionId,
    pub param_info: ParamInfo,
}

impl QuditTensor {
    /// Construct a new tensor object from a tensor expression and param indices.
    pub fn new(indices: Vec<TensorIndex>, expression: ExpressionId, param_info: ParamInfo) -> Self {
        QuditTensor {
            indices,
            expression,
            param_info,
        }
    }

    pub fn num_indices(&self) -> usize {
        self.indices.len()
    }

    pub fn shape(&self) -> GenerationShape {
        (&self.indices).into()
    }

    /// Returns a vector of index IDs for all batch legs of the tensor.
    pub fn batch_indices(&self) -> Vec<usize> {
        self.indices
            .iter()
            .filter_map(|index| match index.direction() {
                IndexDirection::Batch => Some(index.index_id()),
                _ => None,
            })
            .collect()
    }

    /// Returns a vector of index IDs for all output legs of the tensor.
    pub fn output_indices(&self) -> Vec<usize> {
        self.indices
            .iter()
            .filter_map(|index| match index.direction() {
                IndexDirection::Output => Some(index.index_id()),
                _ => None,
            })
            .collect()
    }

    /// Returns a vector of index IDs for all input legs of the tensor.
    pub fn input_indices(&self) -> Vec<usize> {
        self.indices
            .iter()
            .filter_map(|index| match index.direction() {
                IndexDirection::Input => Some(index.index_id()),
                _ => None,
            })
            .collect()
    }

    /// Returns a vector of sizes for all batch legs of the tensor.
    pub fn batch_sizes(&self) -> Vec<usize> {
        self.indices
            .iter()
            .filter_map(|index| match index.direction() {
                IndexDirection::Batch => Some(index.index_size()),
                _ => None,
            })
            .collect()
    }

    /// Returns a vector of sizes for all output legs of the tensor.
    pub fn output_sizes(&self) -> Vec<usize> {
        self.indices
            .iter()
            .filter_map(|index| match index.direction() {
                IndexDirection::Output => Some(index.index_size()),
                _ => None,
            })
            .collect()
    }

    /// Returns a vector of sizes for all input legs of the tensor.
    pub fn input_sizes(&self) -> Vec<usize> {
        self.indices
            .iter()
            .filter_map(|index| match index.direction() {
                IndexDirection::Input => Some(index.index_size()),
                _ => None,
            })
            .collect()
    }

    /// Returns the rank of the tensor.
    pub fn rank(&self) -> usize {
        self.indices.len()
    }
}