1use pyo3::prelude::*;
8use scirs2_neural::activations_minimal::{Activation, ReLU, Sigmoid, Softmax, Tanh, GELU};
9use scirs2_numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods};
10
11#[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 fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
42 apply_activation(&self.inner, py, input)
43 }
44
45 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#[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 fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
86 apply_activation(&self.inner, py, input)
87 }
88
89 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#[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 fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
121 apply_activation(&self.inner, py, input)
122 }
123
124 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#[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 fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
159 apply_activation(&self.inner, py, input)
160 }
161
162 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#[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 fn forward(&self, py: Python, input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
197 apply_activation(&self.inner, py, input)
198 }
199
200 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
211fn apply_activation<A: Activation<f64>>(
217 activation: &A,
218 py: Python,
219 input: &Bound<'_, PyAny>,
220) -> PyResult<Py<PyAny>> {
221 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 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
264fn 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 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 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
330pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
335 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 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}