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

use std::any::Any;
use std::rc::Rc;

use crate::*;
use crate::autograd::PermuteGrad;


impl Tensor {
  fn reshape(&mut self, new_shape: Vec<usize>) {
    self.tensor_mut().set_shape(new_shape);
  }

  fn transpose(&self) -> Self {
    // Transpose by swapping dimensions & strides

    let new_storage = self.tensor().transpose();
    let requires_grad = *self.requires_grad();
    let mut result = Tensor::new(new_storage, self.device(), *self.requires_grad());
    
    if requires_grad {
      result.set_grad_fn(Some(Rc::new(PermuteGrad::new(self, &result))));
    }
    
    result
  }

  fn permute(&mut self, dims: &[usize]) {
    self.tensor_mut().permute(dims);
  }

  fn flatten(&mut self) {
    self.tensor_mut().flatten();
  }

  fn squeeze(&mut self) {
    self.tensor_mut().squeeze();
  } 

  fn unsqueeze(&mut self, dim: usize) {
    self.tensor_mut().unsqueeze(dim);
  }

  fn broadcast(&self, new_shape: &[usize]) -> Self {
    let new_storage = self.tensor().broadcast(new_shape);
    
    // When broadcasting, we need to maintain the gradient tracking
    let requires_grad = *self.requires_grad();
    let mut result = Tensor::new(new_storage, self.device(), *self.requires_grad());
    
    // If original tensor requires gradient, the broadcasted tensor
    // should have the same gradient function
    if requires_grad {
      result.set_grad_fn(self.grad_fn());
    }
    
    result
  }

  fn compute_broadcast_shape(&self, target_shape: &[usize]) -> Vec<usize> {
    self.tensor().compute_broadcast_shape(target_shape)
  }

  fn compute_broadcast_strides(&self, target_shape: &[usize]) -> Vec<usize> {
    self.tensor().compute_broadcast_strides(target_shape)
  }

  fn pad_shape(&self, target_rank: usize) -> Vec<usize> {
    self.tensor().pad_shape(target_rank)
  }


  fn broadcast_tensors(a: &Self, b: &Self) -> (Self, Self) {
    if (a.device() != b.device()) { panic!("Tensors not on same device!") }

    let (broadcast_a, broadcast_b) = Storage::broadcast_tensors(a.tensor(), b.tensor());

    let mut tensor_a = Tensor::new(broadcast_a, a.device(), *a.requires_grad());
    let mut tensor_b = Tensor::new(broadcast_b, b.device(), *b.requires_grad());

    if *a.requires_grad() {
      tensor_a.set_grad_fn(a.grad_fn());
    }
    if *b.requires_grad() {
      tensor_b.set_grad_fn(b.grad_fn());
    }

    (tensor_a, tensor_b)
  }
}