tang 0.2.1

Math library for physical reality — geometry, spatial algebra, tensor, training, GPU compute, and 3D gaussian splatting
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

//! Shape-tracked buffer wrapper for ComputeDevice buffers.

use super::device::{ComputeBuffer, ComputeDevice};

/// A tensor: a device buffer paired with shape metadata.
pub struct ComputeTensor<B: ComputeBuffer> {
    pub buffer: B,
    shape: Vec<usize>,
}

impl<B: ComputeBuffer> ComputeTensor<B> {
    /// Upload f32 data to device with the given shape.
    pub fn from_data<D: ComputeDevice<Buffer = B>>(dev: &D, data: &[f32], shape: &[usize]) -> Self {
        let numel: usize = shape.iter().product();
        assert_eq!(
            data.len(),
            numel,
            "data length {} != shape product {}",
            data.len(),
            numel
        );
        Self {
            buffer: dev.upload(data),
            shape: shape.to_vec(),
        }
    }

    /// Create a zero-filled tensor on device.
    pub fn zeros<D: ComputeDevice<Buffer = B>>(dev: &D, shape: &[usize]) -> Self {
        let numel: usize = shape.iter().product();
        let data = vec![0.0f32; numel];
        Self {
            buffer: dev.upload(&data),
            shape: shape.to_vec(),
        }
    }

    /// Wrap an existing buffer with shape metadata.
    pub fn from_buffer(buffer: B, shape: Vec<usize>) -> Self {
        let numel: usize = shape.iter().product();
        assert_eq!(
            buffer.len(),
            numel,
            "buffer len {} != shape product {}",
            buffer.len(),
            numel
        );
        Self { buffer, shape }
    }

    /// Zero-copy reshape. Panics if numel doesn't match.
    pub fn reshape(self, new_shape: &[usize]) -> Self {
        let new_numel: usize = new_shape.iter().product();
        assert_eq!(
            self.numel(),
            new_numel,
            "reshape: {} != {}",
            self.numel(),
            new_numel
        );
        Self {
            buffer: self.buffer,
            shape: new_shape.to_vec(),
        }
    }

    /// Shape of this tensor.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// Total number of elements.
    pub fn numel(&self) -> usize {
        self.shape.iter().product()
    }

    /// Download contents to CPU.
    pub fn to_vec(&self) -> Vec<f32> {
        self.buffer.to_vec()
    }

    /// Transpose a 2D tensor on device.
    pub fn transpose_2d<D: ComputeDevice<Buffer = B>>(self, dev: &D) -> Self {
        assert_eq!(self.shape.len(), 2, "transpose_2d requires 2D tensor");
        let rows = self.shape[0];
        let cols = self.shape[1];
        let buf = dev.transpose_2d(&self.buffer, rows, cols);
        Self::from_buffer(buf, vec![cols, rows])
    }
}