numr 0.5.2

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
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

//! Shape operations on dual tensors

use crate::autograd::DualTensor;
use crate::error::Result;
use crate::runtime::Runtime;

/// Dual reshape: reshape(a, ȧ) = (reshape(a), reshape(ȧ))
pub fn dual_reshape<R>(a: &DualTensor<R>, shape: &[usize]) -> Result<DualTensor<R>>
where
    R: Runtime,
{
    let primal = a.primal().reshape(shape)?;

    let tangent = match a.tangent() {
        Some(at) => Some(at.reshape(shape)?),
        None => None,
    };

    Ok(DualTensor::new(primal, tangent))
}

/// Dual transpose: transpose(a, ȧ) = (transpose(a), transpose(ȧ))
pub fn dual_transpose<R>(a: &DualTensor<R>, dim0: isize, dim1: isize) -> Result<DualTensor<R>>
where
    R: Runtime,
{
    let primal = a.primal().transpose(dim0, dim1)?;

    let tangent = match a.tangent() {
        Some(at) => Some(at.transpose(dim0, dim1)?),
        None => None,
    };

    Ok(DualTensor::new(primal, tangent))
}

/// Dual narrow (slice): narrow(a, ȧ) = (narrow(a), narrow(ȧ))
pub fn dual_narrow<R>(
    a: &DualTensor<R>,
    dim: isize,
    start: usize,
    length: usize,
) -> Result<DualTensor<R>>
where
    R: Runtime,
{
    let primal = a.primal().narrow(dim, start, length)?;

    let tangent = match a.tangent() {
        Some(at) => Some(at.narrow(dim, start, length)?),
        None => None,
    };

    Ok(DualTensor::new(primal, tangent))
}

/// Dual unsqueeze: unsqueeze(a, ȧ) = (unsqueeze(a), unsqueeze(ȧ))
pub fn dual_unsqueeze<R>(a: &DualTensor<R>, dim: isize) -> Result<DualTensor<R>>
where
    R: Runtime,
{
    let primal = a.primal().unsqueeze(dim)?;

    let tangent = match a.tangent() {
        Some(at) => Some(at.unsqueeze(dim)?),
        None => None,
    };

    Ok(DualTensor::new(primal, tangent))
}

/// Dual squeeze: squeeze(a, ȧ) = (squeeze(a), squeeze(ȧ))
pub fn dual_squeeze<R>(a: &DualTensor<R>, dim: Option<isize>) -> DualTensor<R>
where
    R: Runtime,
{
    let primal = a.primal().squeeze(dim);
    let tangent = a.tangent().map(|at| at.squeeze(dim));
    DualTensor::new(primal, tangent)
}

/// Dual permute: permute(a, ȧ) = (permute(a), permute(ȧ))
pub fn dual_permute<R>(a: &DualTensor<R>, dims: &[usize]) -> Result<DualTensor<R>>
where
    R: Runtime,
{
    let primal = a.primal().permute(dims)?;

    let tangent = match a.tangent() {
        Some(at) => Some(at.permute(dims)?),
        None => None,
    };

    Ok(DualTensor::new(primal, tangent))
}

/// Dual contiguous: contiguous(a, ȧ) = (contiguous(a), contiguous(ȧ))
pub fn dual_contiguous<R>(a: &DualTensor<R>) -> DualTensor<R>
where
    R: Runtime,
{
    let primal = a.primal().contiguous();
    let tangent = a.tangent().map(|at| at.contiguous());
    DualTensor::new(primal, tangent)
}