qudit-tensor 0.3.4

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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use qudit_core::ComplexScalar;
use qudit_core::accel::fused_reshape_permute_reshape_into_impl;
use qudit_core::accel::tensor_fused_reshape_permute_reshape_into_prepare;
use qudit_core::memory::MemoryBuffer;
use qudit_expr::{DifferentiationLevel, FUNCTION, GRADIENT, HESSIAN};

use super::super::buffer::SizedTensorBuffer;

pub struct FRPRStruct<C: ComplexScalar, const D: DifferentiationLevel> {
    pub ins: [Vec<isize>; D],
    pub outs: [Vec<isize>; D],
    pub dims: [Vec<usize>; D],
    pub input: SizedTensorBuffer<C>,
    pub output: SizedTensorBuffer<C>,
}

impl<C: ComplexScalar, const D: DifferentiationLevel> FRPRStruct<C, D> {
    pub fn new(
        input: SizedTensorBuffer<C>,
        shape: &[usize],
        perm: &[usize],
        output: SizedTensorBuffer<C>,
    ) -> Self {
        assert!(input.nparams() == output.nparams());

        let mut all_ins = vec![];
        let mut all_outs = vec![];
        let mut all_dims = vec![];

        if D >= FUNCTION {
            let (ins, outs, dims) = tensor_fused_reshape_permute_reshape_into_prepare(
                &input.shape().to_vec(),
                &input
                    .strides()
                    .iter()
                    .map(|&s| s as isize)
                    .collect::<Vec<isize>>(),
                &output.shape().to_vec(),
                &output
                    .strides()
                    .iter()
                    .map(|&s| s as isize)
                    .collect::<Vec<isize>>(),
                shape,
                perm,
            );

            all_ins.push(ins);
            all_outs.push(outs);
            all_dims.push(dims);
        }

        if D >= GRADIENT {
            // Cool thing: gradient tensors are just tensors, and can have a single FRPR be
            // computed for the transpose operation, rather than a for loop
            // We are really just moving the for loop inside the frpr impl call but this
            // gives prepare a chance to optimize it.
            let (ins, outs, dims) = tensor_fused_reshape_permute_reshape_into_prepare(
                &input.shape().gradient_shape(input.nparams() + 1).to_vec(),
                &input
                    .grad_strides()
                    .iter()
                    .map(|&s| s as isize)
                    .collect::<Vec<isize>>(),
                &output.shape().gradient_shape(input.nparams() + 1).to_vec(),
                &output
                    .grad_strides()
                    .iter()
                    .map(|&s| s as isize)
                    .collect::<Vec<isize>>(),
                &([1 + input.nparams()]
                    .iter()
                    .chain(shape.iter())
                    .copied()
                    .collect::<Vec<usize>>()),
                &([0]
                    .into_iter()
                    .chain(perm.iter().map(|&p| p + 1))
                    .collect::<Vec<usize>>()),
            );

            all_ins.push(ins);
            all_outs.push(outs);
            all_dims.push(dims);
        }

        if D >= HESSIAN {
            let (ins, outs, dims) = tensor_fused_reshape_permute_reshape_into_prepare(
                &(input.shape().gradient_shape(input.nparams() + 1)
                    + input.shape().hessian_shape(input.nparams()))
                .to_vec(),
                &input
                    .grad_strides()
                    .iter()
                    .map(|&s| s as isize)
                    .collect::<Vec<isize>>(),
                &(output.shape().gradient_shape(output.nparams() + 1)
                    + output.shape().hessian_shape(output.nparams()))
                .to_vec(),
                &output
                    .grad_strides()
                    .iter()
                    .map(|&s| s as isize)
                    .collect::<Vec<isize>>(),
                &([1 + input.nparams() + (input.nparams() * (input.nparams() + 1) / 2)]
                    .iter()
                    .chain(shape.iter())
                    .copied()
                    .collect::<Vec<usize>>()),
                &([0]
                    .into_iter()
                    .chain(perm.iter().map(|&p| p + 1))
                    .collect::<Vec<usize>>()),
            );

            all_ins.push(ins);
            all_outs.push(outs);
            all_dims.push(dims);
        }

        let ins = all_ins
            .try_into()
            .expect("Failed to calculate input strides.");
        let outs = all_outs
            .try_into()
            .expect("Failed to calculate output strides.");
        let dims = all_dims
            .try_into()
            .expect("Failed to calculate tensor dimensions.");

        Self {
            ins,
            outs,
            dims,
            input,
            output,
        }
    }

    #[inline(always)]
    pub unsafe fn evaluate<const E: DifferentiationLevel>(&self, memory: &mut MemoryBuffer<C>) {
        unsafe {
            fused_reshape_permute_reshape_into_impl(
                self.input.as_ptr(memory),
                self.output.as_ptr_mut(memory),
                &self.ins[E - 1],
                &self.outs[E - 1],
                &self.dims[E - 1],
            );
        }
    }
}