acyclib 0.3.0

ML library for directed acyclic tensor graphs.
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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

pub mod format;
pub mod manager;
pub mod node;
pub mod topo;

use std::{
    collections::{HashMap, HashSet},
    fmt::Debug,
};

pub use node::{Node, NodeId};

#[derive(Clone, Debug)]
pub enum DAGraphError {
    NodeDoesNotExist,
    NodeWithIdAlreadyExists,
    NodeIsNotRoot,
    FailedTypeCheck,
    Cyclic,
    Message(String),
}

pub trait Operation<Ty: Clone + PartialEq>: Clone {
    fn parents(&self) -> HashSet<NodeId>;

    fn out_type(&self, graph: &DAGraph<Ty, Self>) -> Result<Ty, DAGraphError>;
}

#[derive(Clone)]
pub struct DAGraph<Ty: Clone + PartialEq, Op: Operation<Ty>> {
    nodes: HashMap<NodeId, Node<Ty, Op>>,
}

impl<Ty: Clone + PartialEq, Op: Operation<Ty>> Default for DAGraph<Ty, Op> {
    fn default() -> Self {
        Self { nodes: HashMap::new() }
    }
}

impl<Ty: Clone + PartialEq, Op: Operation<Ty>> DAGraph<Ty, Op> {
    pub fn get(&self, id: NodeId) -> Result<&Node<Ty, Op>, DAGraphError> {
        self.nodes.get(&id).ok_or(DAGraphError::NodeDoesNotExist)
    }

    fn get_mut(&mut self, id: NodeId) -> Result<&mut Node<Ty, Op>, DAGraphError> {
        self.nodes.get_mut(&id).ok_or(DAGraphError::NodeDoesNotExist)
    }

    pub fn roots(&self) -> HashSet<NodeId> {
        self.nodes.values().filter_map(|node| (node.children == 0).then_some(node.id)).collect()
    }

    pub fn topo_order(&self) -> Result<Vec<NodeId>, DAGraphError> {
        let edges_rev =
            self.nodes.iter().map(|(&idx, data)| (idx.0, data.op.parents().iter().map(|x| x.0).collect())).collect();

        topo::topo_order(edges_rev).ok_or(DAGraphError::Cyclic).map(|x| x.into_iter().map(NodeId).collect())
    }

    pub fn add_node(&mut self, op: impl Into<Op>) -> Result<NodeId, DAGraphError> {
        let op = op.into();
        let ty = op.out_type(self)?;
        let node = Node::<Ty, Op>::new(ty, op);
        let id = node.id;

        for parent in node.op.parents() {
            self.get_mut(parent)?.children += 1;
        }

        self.nodes.insert(id, node).map_or(Ok(()), |_| Err(DAGraphError::NodeWithIdAlreadyExists))?;

        Ok(id)
    }

    pub fn remove(&mut self, id: NodeId) -> Result<(), DAGraphError> {
        let node = self.get(id)?;

        if node.children > 0 {
            return Err(DAGraphError::NodeIsNotRoot);
        }

        for parent in node.op.parents() {
            self.get_mut(parent)?.children -= 1;
        }

        self.nodes.remove(&id).expect("Already verified node is present!");

        Ok(())
    }

    pub fn replace_op(&mut self, id: NodeId, new_op: impl Into<Op>) -> Result<(), DAGraphError> {
        let mut new_op = new_op.into();

        if self.get(id)?.ty != new_op.out_type(self)? {
            return Err(DAGraphError::FailedTypeCheck);
        }

        for parent in new_op.parents() {
            self.get_mut(parent)?.children += 1;
        }

        let node = self.get_mut(id)?;

        std::mem::swap(&mut new_op, &mut node.op);

        for parent in new_op.parents() {
            self.get_mut(parent)?.children -= 1;
        }

        // replacing op is an opportunity to introduce cycles...
        self.topo_order()?;

        Ok(())
    }

    pub fn eliminate_dead_nodes(&mut self, required: HashSet<NodeId>) -> Result<(), DAGraphError> {
        for id in self.topo_order()?.into_iter().rev() {
            if !required.contains(&id) && self.get(id)?.children == 0 {
                self.remove(id)?;
            }
        }

        Ok(())
    }
}