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

use crate::{reduce_grad, tensor::*};
use super::super::grad::*;


#[derive(Debug)]
pub struct PermuteGrad {
  input: Tensor,
  output: Tensor,
}


impl PermuteGrad {
  pub fn new(input: &Tensor, output: &Tensor) -> Self {
    PermuteGrad {
      input: input.clone(),
      output: output.clone(),
    }
  }
}

impl GradientFunction for PermuteGrad {
  fn backward(&self) {
    let out_grad = self.output.grad().unwrap();
    let out_grad = out_grad.borrow();

    // Get input gradient if it exists (it should since we're backpropagating)
    if let Some(input_grad) = &self.input.grad() {
      // Determine the permutation that was applied
      // Compare input and output shapes/strides
      let input_shape = self.input.tensor().shape();
      let output_shape = self.output.tensor().shape();
      
      // Find the permutation by matching dimensions
      let mut permutation: Vec<usize> = Vec::new();
      for i in 0..input_shape.len() {
        for j in 0..output_shape.len() {
          if input_shape[i] == output_shape[j] {
            permutation.push(j);
            break;
          }
        }
      }
      
      // Create inverse permutation array
      let mut inverse_perm = vec![0; permutation.len()];
      for (i, &p) in permutation.iter().enumerate() {
        inverse_perm[p] = i;
      }
      
      // Apply inverse permutation to gradient
      let mut grad_tensor = out_grad.clone();
      grad_tensor.permute(&inverse_perm);
      
      // Accumulate the gradient
      input_grad.borrow_mut().add_tensor_assign(&grad_tensor);
    }
  
  }

  fn prev(&self) -> Vec<&Tensor> {
    vec![&self.input]
  }
}