burn-tensor 0.1.0

This library provides multiple tensor implementations hidden behind an easy to use API that supports reverse mode automatic differentiation. 
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

use crate::graph::{node::BackwardNode, ops::RecordedOpsParentRef};
use std::collections::HashSet;

pub trait GraphTraversal {
    fn traverse<F: FnMut(RecordedOpsParentRef)>(&self, callback: F);
}

#[derive(new)]
pub struct BreadthFirstSearch<'a, T> {
    node: &'a BackwardNode<T>,
}

impl<'a, T> GraphTraversal for BreadthFirstSearch<'a, T> {
    fn traverse<F: FnMut(RecordedOpsParentRef)>(&self, mut callback: F) {
        let mut visited = HashSet::with_capacity(self.node.order);
        let mut parents = Vec::with_capacity(self.node.order);

        visited.insert(self.node.id.clone());
        parents.append(&mut self.node.ops.backward_parents());

        loop {
            let node = match parents.pop() {
                Some(node) => node,
                None => break,
            };

            let id = node.id();
            visited.insert(id.clone());

            for parent in node.backward_parents() {
                let id = parent.id();

                if !visited.contains(id) {
                    parents.push(parent);
                }
            }

            callback(node);
        }
    }
}