Skip to main content

scirs2_python/
neural.rs

1//! Python bindings for scirs2-neural
2//!
3//! This module provides Python bindings for neural network activation functions
4//! and utilities. Full layer-based training requires scirs2-autograd's computational
5//! graph system. For comprehensive neural network training, use PyTorch or TensorFlow.
6
7use pyo3::prelude::*;
8use scirs2_neural::activations_minimal::{Activation, ReLU, Sigmoid, Softmax, Tanh, GELU};
9use scirs2_numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods};
10
11// ============================================================================
12// Activation Function Classes
13// ============================================================================
14
15/// ReLU activation function
16///
17/// Applies: f(x) = max(0, x)
18///
19/// Example:
20///     relu = scirs2.ReLU()
21///     output = relu.forward(input_array)
22#[pyclass(name = "ReLU")]
23pub struct PyReLU {
24    inner: ReLU,
25}
26
27#[pymethods]
28impl PyReLU {
29    #[new]
30    fn new() -> Self {
31        Self { inner: ReLU::new() }
32    }
33
34    /// Forward pass
35    ///
36    /// Args:
37    ///     input (np.ndarray): Input array (any shape)
38    ///
39    /// Returns:
40    ///     np.ndarray: Activated output
41    fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
42        apply_activation(&self.inner, py, input)
43    }
44
45    /// Backward pass (gradient computation)
46    ///
47    /// Args:
48    ///     grad_output (np.ndarray): Gradient from next layer
49    ///     input (np.ndarray): Original input to forward pass
50    ///
51    /// Returns:
52    ///     np.ndarray: Gradient with respect to input
53    fn backward(
54        &self,
55        py: Python,
56        grad_output: &Bound<'_, PyAny>,
57        input: &Bound<'_, PyAny>,
58    ) -> PyResult<Py<PyAny>> {
59        apply_activation_backward(&self.inner, py, grad_output, input)
60    }
61}
62
63/// Sigmoid activation function
64///
65/// Applies: f(x) = 1 / (1 + exp(-x))
66///
67/// Example:
68///     sigmoid = scirs2.Sigmoid()
69///     output = sigmoid.forward(input_array)
70#[pyclass(name = "Sigmoid")]
71pub struct PySigmoid {
72    inner: Sigmoid,
73}
74
75#[pymethods]
76impl PySigmoid {
77    #[new]
78    fn new() -> Self {
79        Self {
80            inner: Sigmoid::new(),
81        }
82    }
83
84    /// Forward pass
85    fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
86        apply_activation(&self.inner, py, input)
87    }
88
89    /// Backward pass
90    fn backward(
91        &self,
92        py: Python,
93        grad_output: &Bound<'_, PyAny>,
94        input: &Bound<'_, PyAny>,
95    ) -> PyResult<Py<PyAny>> {
96        apply_activation_backward(&self.inner, py, grad_output, input)
97    }
98}
99
100/// Tanh activation function
101///
102/// Applies: f(x) = tanh(x)
103///
104/// Example:
105///     tanh = scirs2.Tanh()
106///     output = tanh.forward(input_array)
107#[pyclass(name = "Tanh")]
108pub struct PyTanh {
109    inner: Tanh,
110}
111
112#[pymethods]
113impl PyTanh {
114    #[new]
115    fn new() -> Self {
116        Self { inner: Tanh::new() }
117    }
118
119    /// Forward pass
120    fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
121        apply_activation(&self.inner, py, input)
122    }
123
124    /// Backward pass
125    fn backward(
126        &self,
127        py: Python,
128        grad_output: &Bound<'_, PyAny>,
129        input: &Bound<'_, PyAny>,
130    ) -> PyResult<Py<PyAny>> {
131        apply_activation_backward(&self.inner, py, grad_output, input)
132    }
133}
134
135/// GELU activation function
136///
137/// Gaussian Error Linear Unit activation.
138///
139/// Example:
140///     gelu = scirs2.GELU()
141///     output = gelu.forward(input_array)
142#[pyclass(name = "GELU")]
143pub struct PyGELU {
144    inner: GELU,
145}
146
147#[pymethods]
148impl PyGELU {
149    #[new]
150    #[pyo3(signature = (fast=false))]
151    fn new(fast: bool) -> Self {
152        Self {
153            inner: if fast { GELU::fast() } else { GELU::new() },
154        }
155    }
156
157    /// Forward pass
158    fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
159        apply_activation(&self.inner, py, input)
160    }
161
162    /// Backward pass
163    fn backward(
164        &self,
165        py: Python,
166        grad_output: &Bound<'_, PyAny>,
167        input: &Bound<'_, PyAny>,
168    ) -> PyResult<Py<PyAny>> {
169        apply_activation_backward(&self.inner, py, grad_output, input)
170    }
171}
172
173/// Softmax activation function
174///
175/// Applies: f(x)_i = exp(x_i) / sum(exp(x_j))
176///
177/// Example:
178///     softmax = scirs2.Softmax(axis=-1)
179///     output = softmax.forward(input_array)
180#[pyclass(name = "Softmax")]
181pub struct PySoftmax {
182    inner: Softmax,
183}
184
185#[pymethods]
186impl PySoftmax {
187    #[new]
188    #[pyo3(signature = (axis=-1))]
189    fn new(axis: isize) -> Self {
190        Self {
191            inner: Softmax::new(axis),
192        }
193    }
194
195    /// Forward pass
196    fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
197        apply_activation(&self.inner, py, input)
198    }
199
200    /// Backward pass
201    fn backward(
202        &self,
203        py: Python,
204        grad_output: &Bound<'_, PyAny>,
205        input: &Bound<'_, PyAny>,
206    ) -> PyResult<Py<PyAny>> {
207        apply_activation_backward(&self.inner, py, grad_output, input)
208    }
209}
210
211// ============================================================================
212// Helper Functions
213// ============================================================================
214
215/// Apply activation function to NumPy array
216fn apply_activation<A: Activation<f64>>(
217    activation: &A,
218    py: Python,
219    input: &Bound<'_, PyAny>,
220) -> PyResult<Py<PyAny>> {
221    // Try 1D array
222    if let Ok(arr1d) = input.cast::<PyArray1<f64>>() {
223        let binding = arr1d.readonly();
224        let data = binding.as_array().to_owned();
225        let dyn_input = data.into_dyn();
226
227        let output = activation.forward(&dyn_input).map_err(|e| {
228            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Activation error: {}", e))
229        })?;
230
231        let out1d = output
232            .into_dimensionality::<scirs2_core::ndarray::Ix1>()
233            .map_err(|e| {
234                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Dimension error: {}", e))
235            })?;
236
237        return Ok(out1d.into_pyarray(py).unbind().into());
238    }
239
240    // Try 2D array
241    if let Ok(arr2d) = input.cast::<PyArray2<f64>>() {
242        let binding = arr2d.readonly();
243        let data = binding.as_array().to_owned();
244        let dyn_input = data.into_dyn();
245
246        let output = activation.forward(&dyn_input).map_err(|e| {
247            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Activation error: {}", e))
248        })?;
249
250        let out2d = output
251            .into_dimensionality::<scirs2_core::ndarray::Ix2>()
252            .map_err(|e| {
253                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Dimension error: {}", e))
254            })?;
255
256        return Ok(out2d.into_pyarray(py).unbind().into());
257    }
258
259    Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
260        "Input must be 1D or 2D float64 numpy array",
261    ))
262}
263
264/// Apply activation backward pass
265fn apply_activation_backward<A: Activation<f64>>(
266    activation: &A,
267    py: Python,
268    grad_output: &Bound<'_, PyAny>,
269    input: &Bound<'_, PyAny>,
270) -> PyResult<Py<PyAny>> {
271    // Try 1D arrays
272    if let (Ok(grad1d), Ok(inp1d)) = (
273        grad_output.cast::<PyArray1<f64>>(),
274        input.cast::<PyArray1<f64>>(),
275    ) {
276        let grad_binding = grad1d.readonly();
277        let grad_data = grad_binding.as_array().to_owned().into_dyn();
278
279        let inp_binding = inp1d.readonly();
280        let inp_data = inp_binding.as_array().to_owned().into_dyn();
281
282        let grad_input = activation.backward(&grad_data, &inp_data).map_err(|e| {
283            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
284                "Activation backward error: {}",
285                e
286            ))
287        })?;
288
289        let out1d = grad_input
290            .into_dimensionality::<scirs2_core::ndarray::Ix1>()
291            .map_err(|e| {
292                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Dimension error: {}", e))
293            })?;
294
295        return Ok(out1d.into_pyarray(py).unbind().into());
296    }
297
298    // Try 2D arrays
299    if let (Ok(grad2d), Ok(inp2d)) = (
300        grad_output.cast::<PyArray2<f64>>(),
301        input.cast::<PyArray2<f64>>(),
302    ) {
303        let grad_binding = grad2d.readonly();
304        let grad_data = grad_binding.as_array().to_owned().into_dyn();
305
306        let inp_binding = inp2d.readonly();
307        let inp_data = inp_binding.as_array().to_owned().into_dyn();
308
309        let grad_input = activation.backward(&grad_data, &inp_data).map_err(|e| {
310            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
311                "Activation backward error: {}",
312                e
313            ))
314        })?;
315
316        let out2d = grad_input
317            .into_dimensionality::<scirs2_core::ndarray::Ix2>()
318            .map_err(|e| {
319                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Dimension error: {}", e))
320            })?;
321
322        return Ok(out2d.into_pyarray(py).unbind().into());
323    }
324
325    Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
326        "Inputs must be 1D or 2D float64 numpy arrays",
327    ))
328}
329
330// ============================================================================
331// Module Registration
332// ============================================================================
333
334pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
335    // Register activation function classes
336    m.add_class::<PyReLU>()?;
337    m.add_class::<PySigmoid>()?;
338    m.add_class::<PyTanh>()?;
339    m.add_class::<PyGELU>()?;
340    m.add_class::<PySoftmax>()?;
341
342    // Add module documentation
343    m.add(
344        "__doc__",
345        "Neural network activation functions and utilities\n\n\
346        This module provides standalone activation functions that can be used\n\
347        with NumPy arrays for neural network inference and custom training loops.\n\n\
348        Available activations:\n\
349        - ReLU: Rectified Linear Unit\n\
350        - Sigmoid: Logistic sigmoid\n\
351        - Tanh: Hyperbolic tangent\n\
352        - GELU: Gaussian Error Linear Unit\n\
353        - Softmax: Softmax normalization\n\n\
354        Each activation provides:\n\
355        - forward(input): Forward pass\n\
356        - backward(grad_output, input): Backward pass for gradient computation\n\n\
357        For comprehensive neural network training with automatic differentiation,\n\
358        we recommend using PyTorch or TensorFlow, which integrate seamlessly\n\
359        with scirs2 via NumPy array compatibility.",
360    )?;
361
362    Ok(())
363}