scirs2-autograd 0.3.2

Automatic differentiation module for SciRS2 (scirs2-autograd)
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

use crate::op::{ComputeContext, GradientContext, Op, OpError};
use crate::tensor::Tensor;
use crate::Float;

/// Scalar multiplication operation
pub struct ScalarMulOp<F: Float> {
    pub scalar: F,
}

impl<F: Float> Op<F> for ScalarMulOp<F> {
    fn compute(&self, ctx: &mut ComputeContext<F>) -> Result<(), OpError> {
        let input = ctx.input(0);
        let input_array = input.view();
        let output = input_array.mapv(|x| x * self.scalar);
        ctx.append_output(output);
        Ok(())
    }

    fn grad(&self, ctx: &mut GradientContext<F>) {
        let grad_output = ctx.output_grad();
        let g = ctx.graph();

        // Evaluate gradient tensor to array
        let grad_output_array = match grad_output.eval(g) {
            Ok(arr) => arr,
            Err(_) => {
                ctx.append_input_grad(0, None);
                return;
            }
        };

        // Compute gradient (scalar multiplication)
        let grad_input = grad_output_array.mapv(|x| x * self.scalar);

        // Convert to tensor and append
        let grad_tensor = crate::tensor_ops::convert_to_tensor(grad_input, g);
        ctx.append_input_grad(0, Some(grad_tensor));
    }
}

#[allow(dead_code)]
pub fn scalar_mul<'g, F: Float>(tensor: &Tensor<'g, F>, scalar: F) -> Tensor<'g, F> {
    let g = tensor.graph();
    Tensor::builder(g)
        .append_input(tensor, false)
        .setshape(&crate::tensor_ops::shape(tensor))
        .build(ScalarMulOp { scalar })
}