torsh-python 0.1.2

Python bindings for ToRSh - PyTorch-compatible deep learning in 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
124
125
126
127
128
129
130
131
132
133
134
135
136

//! Activation functions for neural networks

use super::module::PyModule;
use crate::{error::PyResult, py_result, tensor::PyTensor};
use pyo3::prelude::*;

/// ReLU activation function
#[pyclass(name = "ReLU", extends = PyModule)]
pub struct PyReLU {
    inplace: bool,
    training: bool,
}

#[pymethods]
impl PyReLU {
    #[new]
    fn new(inplace: Option<bool>) -> (Self, PyModule) {
        (
            Self {
                inplace: inplace.unwrap_or(false),
                training: true,
            },
            PyModule::new(),
        )
    }

    /// Forward pass through ReLU
    fn forward(&self, input: &PyTensor) -> PyResult<PyTensor> {
        let result = py_result!(input.tensor.relu())?;
        Ok(PyTensor { tensor: result })
    }

    /// String representation
    fn __repr__(&self) -> String {
        if self.inplace {
            "ReLU(inplace=True)".to_string()
        } else {
            "ReLU()".to_string()
        }
    }

    /// Set training mode
    fn train(&mut self, mode: Option<bool>) {
        self.training = mode.unwrap_or(true);
    }

    /// Set evaluation mode
    fn eval(&mut self) {
        self.training = false;
    }

    /// Check if module is in training mode
    fn training(&self) -> bool {
        self.training
    }
}

/// Sigmoid activation function
#[pyclass(name = "Sigmoid", extends = PyModule)]
pub struct PySigmoid {
    training: bool,
}

#[pymethods]
impl PySigmoid {
    #[new]
    fn new() -> (Self, PyModule) {
        (Self { training: true }, PyModule::new())
    }

    /// Forward pass through Sigmoid
    fn forward(&self, input: &PyTensor) -> PyResult<PyTensor> {
        let result = py_result!(input.tensor.sigmoid())?;
        Ok(PyTensor { tensor: result })
    }

    /// String representation
    fn __repr__(&self) -> String {
        "Sigmoid()".to_string()
    }

    /// Set training mode
    fn train(&mut self, mode: Option<bool>) {
        self.training = mode.unwrap_or(true);
    }

    /// Set evaluation mode
    fn eval(&mut self) {
        self.training = false;
    }

    /// Check if module is in training mode
    fn training(&self) -> bool {
        self.training
    }
}

/// Tanh activation function
#[pyclass(name = "Tanh", extends = PyModule)]
pub struct PyTanh {
    training: bool,
}

#[pymethods]
impl PyTanh {
    #[new]
    fn new() -> (Self, PyModule) {
        (Self { training: true }, PyModule::new())
    }

    /// Forward pass through Tanh
    fn forward(&self, input: &PyTensor) -> PyResult<PyTensor> {
        let result = py_result!(input.tensor.tanh())?;
        Ok(PyTensor { tensor: result })
    }

    /// String representation
    fn __repr__(&self) -> String {
        "Tanh()".to_string()
    }

    /// Set training mode
    fn train(&mut self, mode: Option<bool>) {
        self.training = mode.unwrap_or(true);
    }

    /// Set evaluation mode
    fn eval(&mut self) {
        self.training = false;
    }

    /// Check if module is in training mode
    fn training(&self) -> bool {
        self.training
    }
}