rustyml 0.17.0

//! Permute layer that reorders the axes after the batch axis, and parks the input shape for
//! backpropagation
//!
//! The configured order leaves the batch axis first. The backward pass applies the inverse
//! order to restore the input axes.

use crate::error::Error;
use crate::neural_network::layers::ParamCounts;
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};
use ndarray::IxDyn;

/// Reorders the axes after the batch axis
///
/// The batch axis stays at axis 0 and never moves, so 1 layer instance serves every batch size.
/// `dims` names the new order of the remaining axes, counting from 1.
/// `Permute::new(vec![2, 1])` on a `[batch, height, width]` input gives `[batch, width, height]`
///
/// The input rank must be `dims.len() + 1`. A rank the layer does not serve is an
/// `Error::InvalidInput` at the forward pass, because the constructor sees no tensor
///
/// Unlike [`Reshape`](crate::neural_network::layers::reshape::Reshape), this layer moves data.
/// A reshape reads the same buffer in the same order under a new shape. A permute reads it in a
/// new order, so the output holds the same values at different positions. When the axis meaning
/// must follow the values, use a permute. 1 example is a layer that reads the last axis as
/// the channel axis
///
/// # Notes
///
/// The axis order changes which stride pattern the data occupies, but not the output's layout.
/// Every layer here emits C order, so this layer copies the reordered view into a fresh
/// array before it returns
///
/// # Examples
///
/// ```rust
/// use rustyml::neural_network::Shape;
/// use rustyml::neural_network::sequential::SequentialBuilder;
/// use rustyml::neural_network::layers::*;
/// use rustyml::neural_network::optimizers::*;
/// use rustyml::neural_network::losses::*;
/// use ndarray::Array3;
///
/// // A rank-3 input: 2 samples, 5 steps, 3 features.
/// let x = Array3::ones((2, 5, 3)).into_dyn();
///
/// // Swap the step axis and the feature axis.
/// let mut model = SequentialBuilder::new()
///     .add(Permute::new(vec![2, 1]).unwrap())
///     .build(&Shape::known(x.shape()))
///     .unwrap();
/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
///
/// // View model structure.
/// model.summary();
///
/// let swapped = model.predict(&x).unwrap();
///
/// // The batch axis stays first, and the other 2 axes trade places.
/// assert_eq!(swapped.shape(), &[2, 3, 5]);
/// ```
#[derive(Debug)]
pub struct Permute {
    /// Axis order the forward pass applies, batch axis included and counted from 0
    ///
    /// Element 0 is always 0. The rest is the `dims` argument, which counts from 1
    forward_axes: Vec<usize>,
    /// Axis order the backward pass applies, the inverse of `forward_axes`
    backward_axes: Vec<usize>,
    /// Shape the layer was built for, batch axis first. `None` before the build
    built: Option<Shape>,
}

impl Permute {
    /// Creates a new Permute layer
    ///
    /// # Parameters
    ///
    /// - `dims` - New order of the axes after the batch axis, counting from 1. Do not include
    ///   the batch axis. The entries must be a permutation of `1..=dims.len()`, so each axis
    ///   appears exactly once
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - New `Permute` layer instance
    ///
    /// # Notes
    ///
    /// This constructor rejects an empty `dims`. Such a layer reorders nothing, and every other
    /// layer here needs a batch axis and at least 1 more axis
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If `dims` is empty, or if it is not a permutation of
    ///   `1..=dims.len()`
    pub fn new(dims: Vec<usize>) -> Result<Self, Error> {
        if dims.is_empty() {
            return Err(Error::invalid_parameter(
                "dims",
                "is empty, and a permutation must name at least 1 axis",
            ));
        }

        let mut seen = vec![false; dims.len()];
        for &axis in &dims {
            match axis.checked_sub(1).and_then(|index| seen.get_mut(index)) {
                Some(slot) if !*slot => *slot = true,
                Some(_) => {
                    return Err(Error::invalid_parameter(
                        "dims",
                        format!("names axis {axis} more than once"),
                    ));
                }
                None => {
                    return Err(Error::invalid_parameter(
                        "dims",
                        format!(
                            "holds {}, and every entry must be between 1 and {}",
                            axis,
                            dims.len()
                        ),
                    ));
                }
            }
        }

        // The batch axis is axis 0 and never moves, so a `dims` entry is already the 0-based
        // index of the axis it names
        let forward_axes: Vec<usize> = std::iter::once(0).chain(dims).collect();
        let mut backward_axes = vec![0; forward_axes.len()];
        for (output_axis, &input_axis) in forward_axes.iter().enumerate() {
            backward_axes[input_axis] = output_axis;
        }

        Ok(Permute {
            forward_axes,
            backward_axes,
            built: None,
        })
    }

    /// Shape this layer produces from `input_shape`
    fn permuted_shape(&self, input_shape: &[usize]) -> Vec<usize> {
        self.forward_axes.iter().map(|&a| input_shape[a]).collect()
    }

    /// Checks the rank and the element count of a tensor entering the layer
    ///
    /// # Errors
    ///
    /// - `Error::InvalidInput` - If the rank is not `dims.len() + 1`
    /// - `Error::EmptyInput` - If any axis has an extent of 0
    fn validate(&self, input: &Tensor) -> Result<(), Error> {
        if input.ndim() != self.forward_axes.len() {
            return Err(Error::invalid_input(format!(
                "Permute layer expects a {}D input, got a {}D tensor",
                self.forward_axes.len(),
                input.ndim()
            )));
        }
        if input.is_empty() {
            return Err(Error::empty_input("input tensor"));
        }
        Ok(())
    }
}

/// Copies `input` under a new axis order, into a tensor that is in C order
///
/// `permuted_axes` reorders the strides and shares the buffer, so it moves no data. `to_owned`
/// on such a view keeps those reordered strides, which leaves the result out of C order. Every
/// layer here emits C order, so a consumer can read any layer output as 1 contiguous slice.
/// This layer copies the data to meet that contract
fn permute_into(input: &Tensor, axes: &[usize]) -> Tensor {
    let view = input.view().permuted_axes(IxDyn(axes));
    let mut output = Tensor::zeros(view.raw_dim());
    output.assign(&view);
    output
}

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

    built_layer_shape_functions!();

    no_trainable_parameters_layer_functions!();
}

impl UnaryLayer for Permute {
    /// Records the shape the layer reorders. The layer holds no array, so nothing is allocated.
    /// The shape algebra checks the rank
    fn build(&mut self, input: &Shape) -> Result<(), Error> {
        let Some(built) = start_build(&self.built, "Permute", 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> {
        self.validate(input)?;

        if ctx.is_training() {
            ctx.push_cache("Permute", input.shape().to_vec());
        }

        Ok(permute_into(input, &self.forward_axes))
    }

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

        let expected = self.permuted_shape(&input_shape);
        if grad_output.shape() != expected.as_slice() {
            return Err(Error::shape_mismatch(expected, grad_output.shape()));
        }

        // A permute moves each value to 1 new position, so the gradient runs back through the
        // inverse order
        Ok(permute_into(grad_output, &self.backward_axes))
    }

    /// The layer moves axes and changes no extent, so the output axes are the input axes in
    /// the configured order. Element 0 of that order is always the batch axis
    fn compute_output_shape(&self, input: &Shape) -> Result<Shape, Error> {
        input.check_rank("Permute", self.forward_axes.len())?;
        let axes = input.axes();
        Ok(Shape::new(
            self.forward_axes.iter().map(|&axis| axes[axis]).collect(),
        ))
    }
}