ferrite-dl 0.2.0

Deep learning library written in pure 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
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

use std::rc::Rc;
use std::cell::RefCell;
use super::storage::*;
use crate::{grad_storage, GradientFunction, CpuStorage};
use std::collections::HashSet;


pub type GradientStorage = Rc<RefCell<Storage>>;

#[derive(Clone)]
pub struct Tensor {
  pub storage: Storage,
  device: Device,
  requires_grad: bool,
  grad_fn: Option<Rc<dyn GradientFunction>>,
  grad: Option<GradientStorage>,
}

impl Tensor {
  pub fn new(storage: Storage, device: Device, requires_grad: bool) -> Self {
    let grad = if requires_grad {
      Some(Rc::new(RefCell::new(Storage::zeros(storage.shape().clone(), Some(device), None))))
    } else {
      None
    };
    
    Tensor {
      storage: storage,
      device: device,
      requires_grad: requires_grad,
      grad_fn: None,
      grad: grad,
    }
  }

  pub fn view(&self, tensor: Storage) -> Self {
    Tensor {
      storage: tensor,
      device: self.device,
      requires_grad: self.requires_grad,
      grad_fn: self.grad_fn.clone(),
      grad: self.grad.clone(),
    }
  }

  pub fn tensor(&self) -> &Storage {
    &self.storage
  }

  pub fn tensor_mut(&mut self) -> &mut Storage {
    &mut self.storage
  }

  pub fn device(&self) -> Device {
    self.device
  }

  pub fn requires_grad(&self) -> &bool {
    &self.requires_grad
  }

  pub fn grad_fn(&self) -> Option<Rc<dyn GradientFunction>> {
    self.grad_fn.clone()
  }

  pub fn set_grad_fn(&mut self, grad_fn: Option<Rc<dyn GradientFunction>>) {
    self.grad_fn = grad_fn;
  }

  pub fn grad(&self) -> Option<GradientStorage> {
    self.grad.clone()
  }

  pub fn grad_mut(&mut self) -> GradientStorage {
    self.grad.clone().expect("Grad can't be empty")
  }

  pub fn shape(&self) -> &Vec<usize> {
    &self.tensor().shape()
  }

  pub fn backward(&mut self) {
    // Verify we're starting with a scalar
    if self.tensor().shape().len() != 1 || self.tensor().shape()[0] != 1 {
      panic!("backward() can only be called on scalar tensors");
    }

    // Initialize gradient for final output (always 1.0 for scalar outputs)
    if let Some(grad) = &self.grad {
      grad.borrow_mut().set_data(vec![1.0]);
    } else {
      panic!("Called backward on tensor that doesn't require grad");
    }

    // Build computation graph in topological order
    let mut topo = Vec::new();
    let mut visited = HashSet::new();

    fn build_topo(
      node: &Tensor, 
      topo: &mut Vec<Rc<dyn GradientFunction>>, 
      visited: &mut HashSet<*const dyn GradientFunction>
    ) {
      if let Some(grad_fn) = &node.grad_fn {
        let ptr = Rc::as_ptr(grad_fn) as *const dyn GradientFunction;
        if !visited.contains(&ptr) {
          visited.insert(ptr);
          for parent in grad_fn.prev() {
            build_topo(parent, topo, visited);
          }
          topo.push(grad_fn.clone());
        }
      }
    }

    build_topo(self, &mut topo, &mut visited);

    // Execute backward passes in reverse order
    for grad_fn in topo.iter().rev() {
      grad_fn.backward();
    }
  }
}