ndtensor 0.1.0

ndtensor is an n-dimensional tensor library for Rust.
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

/*
    Appellation: utils <module>
    Contrib: FL03 <jo3mccain@icloud.com>
*/
use crate::prelude::{TensorExpr, TensorId};
use crate::TensorBase;
use nd::{Array, Dimension, IntoDimension, RawData, RawDataClone};
use num::Float;
use std::collections::HashMap;

/// Hashes a dimension using the [DefaultHasher].
#[cfg(feature = "std")]
pub fn hash_dim<D>(dim: impl IntoDimension<Dim = D>) -> u64
where
    D: Dimension,
{
    use std::hash::{DefaultHasher, Hash, Hasher};
    let dim = dim.into_dimension();
    let mut s = DefaultHasher::new();
    for i in dim.slice() {
        i.hash(&mut s);
    }
    s.finish()
}

pub fn linarr<A, D>(dim: impl IntoDimension<Dim = D>) -> Array<A, D>
where
    A: Float,
    D: Dimension,
{
    let dim = dim.into_dimension();
    let dview = dim.as_array_view();
    let n = dview.product();
    Array::linspace(A::zero(), A::from(n).unwrap() - A::one(), n)
        .into_shape(dim)
        .expect("linspace err")
}

pub(crate) fn walk<S>(
    scope: TensorBase<S>,
    nodes: Vec<TensorBase<S>>,
    visited: &mut HashMap<TensorId, bool>,
) -> (bool, Vec<TensorBase<S>>)
where
    S: RawData + RawDataClone,
{
    if let Some(&tg) = visited.get(&scope.id()) {
        return (tg, nodes);
    }
    // track the gradient of the current node
    let mut track = false;
    // recursively call on the children nodes
    let mut nodes = if scope.is_variable() {
        // Do not call recursively on the "leaf" nodes.
        track = true;
        nodes
    } else if let Some(op) = scope.op() {
        match op {
            TensorExpr::Binary { lhs, rhs, .. } => {
                let (tg, nodes) = walk(*lhs.clone(), nodes, visited);
                track |= tg;
                let (tg, nodes) = walk(*rhs.clone(), nodes, visited);
                track |= tg;
                nodes
            }
            TensorExpr::Unary { recv, .. } => {
                let (tg, nodes) = walk(*recv.clone(), nodes, visited);
                track |= tg;
                nodes
            }
            _ => nodes,
        }
    } else {
        nodes
    };
    visited.insert(scope.id(), track);
    if track {
        nodes.push(scope);
    }
    (track, nodes)
}