Tensor Operations

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// A simple row-major Tensor with dynamic rank.
5
// Focus: shape/strides, linear offset, broadcasting-aware elementwise ops.
6
template <typename T>
7
class Tensor {
8
public:
9
    vector<size_t> shape;      // sizes per axis
10
    vector<size_t> strides;    // row-major strides
11
    vector<T> data;            // contiguous storage
12
13
    Tensor() = default;
14
    explicit Tensor(const vector<size_t>& shp) : shape(shp) {
15
        compute_strides();
16
        data.assign(numel(), T{});
17
    }
18
19
    size_t rank() const { return shape.size(); }
20
    size_t numel() const {
21
        return accumulate(shape.begin(), shape.end(), (size_t)1, multiplies<size_t>());
22
    }
23
24
    // Compute row-major strides: s_k = product of later dimensions.
25
    void compute_strides() {
26
        strides.assign(rank(), 1);
27
        for (int k = (int)rank() - 2; k >= 0; --k) {
28
            strides[k] = strides[k+1] * shape[k+1];
29
        }
30
    }
31
32
    // Convert a multi-index to linear offset.
33
    size_t offset(const vector<size_t>& idx) const {
34
        size_t off = 0;
35
        for (size_t k = 0; k < idx.size(); ++k) {
36
            off += idx[k] * strides[k];
37
        }
38
        return off;
39
    }
40
41
    // Safe element accessors using a multi-index.
42
    T get(const vector<size_t>& idx) const { return data[offset(idx)]; }
43
    void set(const vector<size_t>& idx, const T& value) { data[offset(idx)] = value; }
44
45
    // Fill with a lambda f(idx_flat) or f(multi-index)
46
    template <class F>
47
    void fill_flat(F f) {
48
        for (size_t i = 0; i < numel(); ++i) data[i] = f(i);
49
    }
50
51
    // Broadcasting utilities
52
    static vector<size_t> broadcast_shape(vector<size_t> a, vector<size_t> b) {
53
        // Right-align shapes
54
        size_t na = a.size(), nb = b.size();
55
        size_t n = max(na, nb);
56
        a.insert(a.begin(), n - na, 1);
57
        b.insert(b.begin(), n - nb, 1);
58
        vector<size_t> out(n);
59
        for (size_t k = 0; k < n; ++k) {
60
            if (a[k] == b[k] || a[k] == 1) out[k] = b[k];
61
            else if (b[k] == 1) out[k] = a[k];
62
            else throw runtime_error("Shapes not broadcastable");
63
        }
64
        return out;
65
    }
66
67
    // Compute effective strides for broadcasting: stride is 0 if dim==1.
68
    static vector<size_t> broadcast_strides(const vector<size_t>& in_shape,
69
                                            const vector<size_t>& in_strides,
70
                                            const vector<size_t>& out_shape) {
71
        size_t n = out_shape.size();
72
        vector<size_t> a = in_shape;
73
        a.insert(a.begin(), n - a.size(), 1);
74
        vector<size_t> s = in_strides;
75
        s.insert(s.begin(), n - s.size(), 0);
76
        vector<size_t> eff(n);
77
        for (size_t k = 0; k < n; ++k) eff[k] = (a[k] == 1 ? 0 : s[k]);
78
        return eff;
79
    }
80
81
    // Generic broadcasted binary elementwise op: out = f(A, B)
82
    template <class F>
83
    static Tensor apply_binary(const Tensor& A, const Tensor& B, F f) {
84
        vector<size_t> out_shape = broadcast_shape(A.shape, B.shape);
85
        Tensor out(out_shape);
86
        // Precompute effective strides aligning to out_shape
87
        vector<size_t> sA = broadcast_strides(A.shape, A.strides, out_shape);
88
        vector<size_t> sB = broadcast_strides(B.shape, B.strides, out_shape);
89
        // Iterate linearly over out tensor using mixed-radix counter
90
        size_t N = out.numel();
91
        vector<size_t> idx(out_shape.size(), 0);
92
        for (size_t flat = 0; flat < N; ++flat) {
93
            // Compute linear offsets for A and B via effective strides
94
            size_t offA = 0, offB = 0;
95
            for (size_t k = 0; k < idx.size(); ++k) {
96
                offA += idx[k] * sA[k];
97
                offB += idx[k] * sB[k];
98
            }
99
            out.data[flat] = f(A.data[offA], B.data[offB]);
100
            // Increment mixed-radix counter (row-major)
101
            for (int k = (int)idx.size() - 1; k >= 0; --k) {
102
                if (++idx[k] < out_shape[k]) break;
103
                idx[k] = 0;
104
            }
105
        }
106
        return out;
107
    }
108
};
109
110
int main() {
111
    // Example: A (2,3) plus B (1,3) via broadcasting along axis 0.
112
    Tensor<double> A({2,3});
113
    Tensor<double> B({1,3});
114
115
    // Fill A with row i times 10 plus column j, and B with a bias per column.
116
    A.fill_flat([&](size_t t){ return (double)t; }); // simple increasing values
117
    B.set({0,0}, 100.0);
118
    B.set({0,1}, 200.0);
119
    B.set({0,2}, 300.0);
120
121
    auto C = Tensor<double>::apply_binary(A, B, [](double x, double y){ return x + y; });
122
    auto D = Tensor<double>::apply_binary(C, C, [](double x, double y){ return x * y; }); // Hadamard square
123
124
    // Print results to verify broadcasting and elementwise ops
125
    cout << "C shape: (2,3)\n";
126
    for (size_t i = 0; i < 2; ++i) {
127
        for (size_t j = 0; j < 3; ++j) cout << C.get({i,j}) << ' ';
128
        cout << '\n';
129
    }
130
    cout << "\nD = C .* C\n";
131
    for (size_t i = 0; i < 2; ++i) {
132
        for (size_t j = 0; j < 3; ++j) cout << D.get({i,j}) << ' ';
133
        cout << '\n';
134
    }
135
}
136

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
template <typename T>
5
struct Tensor {
6
    vector<size_t> shape, strides; vector<T> data;
7
    Tensor() {}
8
    explicit Tensor(const vector<size_t>& shp): shape(shp) {
9
        strides.assign(shape.size(), 1);
10
        for (int k = (int)shape.size()-2; k >= 0; --k) strides[k] = strides[k+1]*shape[k+1];
11
        data.assign(numel(), T{});
12
    }
13
    size_t numel() const { return accumulate(shape.begin(), shape.end(), (size_t)1, multiplies<size_t>()); }
14
    size_t offset(const vector<size_t>& idx) const { size_t off=0; for (size_t k=0;k<idx.size();++k) off += idx[k]*strides[k]; return off; }
15
    T get(const vector<size_t>& idx) const { return data[offset(idx)]; }
16
    void set(const vector<size_t>& idx, const T& v){ data[offset(idx)] = v; }
17
};
18
19
// Iterate over all indices of a shape, calling fn(index).
20
void for_each_index(const vector<size_t>& shape, function<void(const vector<size_t>&)> fn){
21
    vector<size_t> idx(shape.size(),0);
22
    size_t N = accumulate(shape.begin(), shape.end(), (size_t)1, multiplies<size_t>());
23
    for (size_t flat=0; flat<N; ++flat){
24
        fn(idx);
25
        for (int k=(int)shape.size()-1; k>=0; --k){ if (++idx[k] < shape[k]) break; idx[k]=0; }
26
    }
27
}
28
29
// Contract A along axisA with B along axisB. Returns a new tensor.
30
template <typename T>
31
Tensor<T> contract(const Tensor<T>& A, int axisA, const Tensor<T>& B, int axisB){
32
    if (A.shape[axisA] != B.shape[axisB]) throw runtime_error("Contracted axes must match");
33
    size_t r = A.shape[axisA];
34
    // Build output shape = A.shape without axisA concatenated with B.shape without axisB
35
    vector<size_t> out_shape; out_shape.reserve(A.shape.size()+B.shape.size()-2);
36
    for (size_t i=0;i<A.shape.size();++i) if ((int)i!=axisA) out_shape.push_back(A.shape[i]);
37
    for (size_t j=0;j<B.shape.size();++j) if ((int)j!=axisB) out_shape.push_back(B.shape[j]);
38
    Tensor<T> C(out_shape);
39
40
    // Iterate over all output indices and accumulate over k in [0, r)
41
    for_each_index(out_shape, [&](const vector<size_t>& out_idx){
42
        // Split out_idx into parts corresponding to A (without axisA) and B (without axisB)
43
        vector<size_t> idxA_full(A.shape.size());
44
        vector<size_t> idxB_full(B.shape.size());
45
        // Fill idxA_full by inserting from out_idx
46
        size_t p=0;
47
        for (size_t i=0;i<A.shape.size();++i){
48
            if ((int)i==axisA) continue;
49
            idxA_full[i] = out_idx[p++];
50
        }
51
        // Fill idxB_full from the remaining suffix of out_idx
52
        size_t q=0;
53
        for (size_t j=0;j<B.shape.size();++j){
54
            if ((int)j==axisB) continue;
55
            idxB_full[j] = out_idx[p + q++];
56
        }
57
        // Accumulate over contracted index
58
        T acc = T{};
59
        for (size_t k=0;k<r;++k){
60
            idxA_full[axisA] = k;
61
            idxB_full[axisB] = k;
62
            acc += A.get(idxA_full) * B.get(idxB_full);
63
        }
64
        C.data[C.offset(out_idx)] = acc;
65
    });
66
    return C;
67
}
68
69
int main(){
70
    // Example: matrix multiplication via contraction
71
    Tensor<double> A({2,3}); // (2x3)
72
    Tensor<double> B({3,4}); // (3x4)
73
74
    // Fill A and B with small values
75
    for (size_t i=0;i<2;i++) for (size_t k=0;k<3;k++) A.set({i,k}, (i+1)*(k+1));
76
    for (size_t k=0;k<3;k++) for (size_t j=0;j<4;j++) B.set({k,j}, (k+1)+(j));
77
78
    // Contract A axis 1 with B axis 0 => output shape (2,4)
79
    auto C = contract(A, 1, B, 0);
80
81
    cout << "C (2x4) = A(2x3) * B(3x4)\n";
82
    for (size_t i=0;i<2;i++){ for (size_t j=0;j<4;j++) cout << C.get({i,j}) << ' '; cout << '\n'; }
83
}
84

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
template <typename T>
5
struct Tensor {
6
    vector<size_t> shape, strides; vector<T> data;
7
    Tensor(){}
8
    explicit Tensor(const vector<size_t>& shp): shape(shp){
9
        strides.assign(shape.size(),1);
10
        for (int k=(int)shape.size()-2;k>=0;--k) strides[k]=strides[k+1]*shape[k+1];
11
        data.assign(numel(), T{});
12
    }
13
    size_t numel() const { return accumulate(shape.begin(), shape.end(), (size_t)1, multiplies<size_t>()); }
14
    size_t offset(const vector<size_t>& idx) const { size_t off=0; for (size_t k=0;k<idx.size();++k) off += idx[k]*strides[k]; return off; }
15
};
16
17
template <typename T>
18
Tensor<T> permute(const Tensor<T>& A, const vector<int>& axes){
19
    if (axes.size()!=A.shape.size()) throw runtime_error("axes length mismatch");
20
    // New shape order
21
    vector<size_t> new_shape(axes.size());
22
    for (size_t i=0;i<axes.size();++i) new_shape[i]=A.shape[axes[i]];
23
    Tensor<T> B(new_shape);
24
    // Iterate over all indices of B and map back to A
25
    vector<size_t> idxB(new_shape.size(),0);
26
    size_t N=B.numel();
27
    for (size_t flat=0; flat<N; ++flat){
28
        vector<size_t> idxA(A.shape.size());
29
        for (size_t i=0;i<axes.size();++i) idxA[axes[i]] = idxB[i];
30
        B.data[flat] = A.data[A.offset(idxA)];
31
        for (int k=(int)idxB.size()-1;k>=0;--k){ if (++idxB[k] < new_shape[k]) break; idxB[k]=0; }
32
    }
33
    return B;
34
}
35
36
// Reshape with copy (for simplicity). A view-based reshape would share data.
37
template <typename T>
38
Tensor<T> reshape_copy(const Tensor<T>& A, const vector<size_t>& new_shape){
39
    size_t newN = accumulate(new_shape.begin(), new_shape.end(), (size_t)1, multiplies<size_t>());
40
    if (newN != A.numel()) throw runtime_error("element count mismatch in reshape");
41
    Tensor<T> B(new_shape);
42
    B.data = A.data; // copy assignment; same order preserved
43
    return B;
44
}
45
46
int main(){
47
    Tensor<int> A({2,3,4});
48
    // Fill A with 0..23
49
    for (int i=0;i<24;++i) A.data[i]=i;
50
51
    // Permute axes from (0,1,2) to (2,0,1): shape (4,2,3)
52
    auto P = permute(A, {2,0,1});
53
54
    // Reshape P to (4,6)
55
    auto R = reshape_copy(P, {4,6});
56
57
    cout << "P shape (4,2,3) in row-major order first row of last 2D view by printing first 6 elements after reshape:\n";
58
    for (size_t j=0;j<6;++j) cout << R.data[j] << ' ';
59
    cout << "\nTotal elements: " << R.numel() << "\n";
60
}
61