use crate::error::Error;
use crate::neural_network::layers::ParamCounts;
use crate::neural_network::layers::activation::{Activation, ActivationCache, ActivationInput};
use crate::neural_network::layers::validation::start_build;
use crate::neural_network::layers::{
built_layer_shape_functions, no_trainable_parameters_layer_functions,
};
use crate::neural_network::traits::{LayerBase, UnaryLayer};
use crate::neural_network::{Ctx, Shape, Tensor};
#[derive(Debug)]
pub struct ELU {
pub(super) alpha: f32,
built: Option<Shape>,
}
impl ELU {
pub fn new(alpha: f32) -> Result<Self, Error> {
Activation::ELU { alpha }.validate()?;
Ok(ELU { alpha, built: None })
}
}
impl Default for ELU {
fn default() -> Self {
ELU {
alpha: 1.0,
built: None,
}
}
}
impl LayerBase for ELU {
fn layer_type(&self) -> &str {
"ELU"
}
built_layer_shape_functions!();
no_trainable_parameters_layer_functions!();
}
impl UnaryLayer for ELU {
fn build(&mut self, input: &Shape) -> Result<(), Error> {
let Some(built) = start_build(&self.built, "ELU", input)? else {
return Ok(());
};
self.compute_output_shape(&built)?;
self.built = Some(built);
Ok(())
}
fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
if input.is_empty() {
return Err(Error::empty_input("input tensor"));
}
let output = Activation::ELU { alpha: self.alpha }.forward(input)?;
if ctx.is_training() {
ctx.push_cache(
"ELU",
ActivationCache::new(ActivationInput::Output, output.clone()),
);
}
Ok(output)
}
fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
let cache: ActivationCache = ctx.pop_cache("ELU")?;
Activation::ELU { alpha: self.alpha }.backward(&cache, grad_output)
}
}