rustyml 0.17.0

//! Flatten layer that reshapes a 3D, 4D, or 5D tensor into a 2D tensor, and parks the input
//! shape for backpropagation
//!
//! The batch axis stays first, and all other axes collapse in C order. The backward pass
//! restores the cached input shape.

use crate::error::{Context, Error};
use crate::neural_network::layers::ParamCounts;
use crate::neural_network::layers::validation::{start_build, validate_built_input};
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};
use ndarray::IxDyn;

/// Flattens a 3D, 4D, or 5D tensor into a 2D tensor
///
/// Reshapes inputs from feature extraction layers into a format suitable for dense layers
///
/// The reshape itself is layout-agnostic. It collapses every axis after the batch axis in C
/// order, without regard to which feature lands at which output index. Under the crate's
/// channels-last layout, the channel axis is innermost. The flattened vector then runs position
/// by position, with all channels of 1 position adjacent. Dense weights must use this feature
/// order. A shape check alone cannot detect a different feature order
///
/// Input shapes are `[batch_size, length, features]`, `[batch_size, height, width, channels]`,
/// or `[batch_size, depth, height, width, channels]`. The output shape is always
/// `[batch_size, flattened_features]`, where `flattened_features` is the product of all
/// dimensions except the batch size
///
/// # Examples
///
/// ```rust
/// use rustyml::neural_network::sequential::SequentialBuilder;
/// use rustyml::neural_network::layers::*;
/// use rustyml::neural_network::optimizers::*;
/// use rustyml::neural_network::losses::*;
/// use rustyml::neural_network::Shape;
/// use ndarray::Array4;
///
/// // Create a 4D input tensor: [batch_size, height, width, channels].
/// // Batch size=2, 4x4 pixels, 3 channels.
/// let x = Array4::ones((2, 4, 4, 3)).into_dyn();
///
/// // Build a model containing a Flatten layer.
/// let mut model = SequentialBuilder::new()
///     .add(Flatten::new())
///     .build(&Shape::known(&[2, 4, 4, 3]))
///     .unwrap();
/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
///
/// // View model structure.
/// model.summary();
///
/// // Forward propagation.
/// let flattened = model.predict(&x).unwrap();
///
/// // The output shape should be [2, 48].
/// assert_eq!(flattened.shape(), &[2, 48]);
/// ```
#[derive(Debug, Default)]
pub struct Flatten {
    /// Shape the layer was built for, batch axis first. `None` before the build
    built: Option<Shape>,
}

impl Flatten {
    /// Creates a new Flatten layer
    ///
    /// The layer starts with no input shape. [`UnaryLayer::build`] sets it, and a forward pass
    /// on a layer that a caller drives by hand builds it from the tensor that arrives
    ///
    /// # Returns
    ///
    /// - `Self` - New `Flatten` layer instance
    pub fn new() -> Self {
        Self { built: None }
    }
}

impl LayerBase for Flatten {
    fn layer_type(&self) -> &str {
        "Flatten"
    }

    built_layer_shape_functions!();

    no_trainable_parameters_layer_functions!();
}

impl UnaryLayer for Flatten {
    /// Records the shape the layer folds. The layer holds no array, so nothing is allocated
    fn build(&mut self, input: &Shape) -> Result<(), Error> {
        let Some(built) = start_build(&self.built, "Flatten", input)? else {
            return Ok(());
        };
        if !(3..=5).contains(&built.rank()) {
            return Err(Error::invalid_input(format!(
                "Flatten layer expects 3D, 4D, or 5D input, got {}D tensor",
                built.rank()
            )));
        }
        built.split_batch("Flatten")?;
        self.built = Some(built);
        Ok(())
    }

    fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        validate_built_input(&self.built, "Flatten", input.shape())?;
        let input_shape = input.shape();

        if ctx.is_training() {
            ctx.push_cache("Flatten", input_shape.to_vec());
        }

        let batch_size = input_shape[0];
        let flattened_features: usize = input_shape[1..].iter().product();

        Ok(input
            .to_shape(IxDyn(&[batch_size, flattened_features]))
            .unwrap()
            .to_owned())
    }

    fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        let input_shape: Vec<usize> = ctx.pop_cache("Flatten")?;

        let expected_grad_shape = [input_shape[0], input_shape[1..].iter().product()];
        if grad_output.shape() != expected_grad_shape {
            return Err(Error::shape_mismatch(
                expected_grad_shape,
                grad_output.shape(),
            ));
        }

        let reshaped_grad = grad_output
            .to_shape(IxDyn(input_shape.as_slice()))
            .context("reshape gradient")?
            .to_owned();

        Ok(reshaped_grad)
    }

    /// Every axis after the batch axis folds into 1 feature axis
    fn compute_output_shape(&self, input: &Shape) -> Result<Shape, Error> {
        input.check_min_rank("Flatten", 2)?;
        let (batch, tail) = input.split_batch("Flatten")?;
        Ok(Shape::from_batch(batch, &[tail.iter().product()]))
    }
}