✨ Initial implementation of residual neural networks

src/layer/dense.rs

@@ -122,6 +122,14 @@ impl<F, Act, Reg, R: Rng> NeuraDenseLayerPartial<F, Act, Reg, R> {
     }
 }
 
+impl<F: Float, Act: NeuraDerivable<F>, Reg: NeuraDerivable<F>> NeuraShapedLayer
+    for NeuraDenseLayer<F, Act, Reg>
+{
+    fn output_shape(&self) -> NeuraShape {
+        NeuraShape::Vector(self.weights.shape().0)
+    }
+}
+
 impl<
         F: Float + std::fmt::Debug + 'static,
         Act: NeuraDerivable<F>,
@@ -144,10 +152,6 @@ where
             self.regularization,
         ))
     }
-
-    fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
-        NeuraShape::Vector(constructed.weights.shape().0)
-    }
 }
 
 impl<

src/layer/dropout.rs

@@ -35,6 +35,12 @@ impl<R: Rng> NeuraDropoutLayer<R> {
     }
 }
 
+impl<R: Rng> NeuraShapedLayer for NeuraDropoutLayer<R> {
+    fn output_shape(&self) -> NeuraShape {
+        self.shape
+    }
+}
+
 impl<R: Rng> NeuraPartialLayer for NeuraDropoutLayer<R> {
     type Constructed = NeuraDropoutLayer<R>;
 
@@ -45,10 +51,6 @@ impl<R: Rng> NeuraPartialLayer for NeuraDropoutLayer<R> {
         self.mask = DVector::from_element(input_shape.size(), false);
         Ok(self)
     }
-
-    fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
-        constructed.shape
-    }
 }
 
 impl<R: Rng, F: Float> NeuraLayer<DVector<F>> for NeuraDropoutLayer<R> {

src/layer/mod.rs

@@ -48,13 +48,15 @@ impl<Input: Clone> NeuraLayer<Input> for () {
     }
 }
 
+pub trait NeuraShapedLayer {
+    fn output_shape(&self) -> NeuraShape;
+}
+
 pub trait NeuraPartialLayer {
-    type Constructed;
+    type Constructed: NeuraShapedLayer;
     type Err;
 
     fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err>;
-
-    fn output_shape(constructed: &Self::Constructed) -> NeuraShape;
 }
 
 pub trait NeuraTrainableLayerBase<Input>: NeuraLayer<Input> {

src/layer/normalize.rs

@@ -23,6 +23,12 @@ impl NeuraNormalizeLayer {
     }
 }
 
+impl NeuraShapedLayer for NeuraNormalizeLayer {
+    fn output_shape(&self) -> NeuraShape {
+        self.shape
+    }
+}
+
 impl NeuraPartialLayer for NeuraNormalizeLayer {
     type Constructed = NeuraNormalizeLayer;
 
@@ -31,10 +37,6 @@ impl NeuraPartialLayer for NeuraNormalizeLayer {
     fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
         Ok(Self { shape: input_shape })
     }
-
-    fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
-        constructed.shape
-    }
 }
 
 impl<F: Float + Scalar> NeuraLayer<DVector<F>> for NeuraNormalizeLayer {

src/layer/softmax.rs

@@ -41,6 +41,12 @@ impl<F: Float + Scalar + NumAssignOps> NeuraLayer<DVector<F>> for NeuraSoftmaxLa
     }
 }
 
+impl NeuraShapedLayer for NeuraSoftmaxLayer {
+    fn output_shape(&self) -> NeuraShape {
+        self.shape
+    }
+}
+
 impl NeuraPartialLayer for NeuraSoftmaxLayer {
     type Constructed = Self;
     type Err = ();
@@ -48,10 +54,6 @@ impl NeuraPartialLayer for NeuraSoftmaxLayer {
     fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
         Ok(Self { shape: input_shape })
     }
-
-    fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
-        constructed.shape
-    }
 }
 
 impl<F: Float + Scalar + NumAssignOps> NeuraTrainableLayerBase<DVector<F>> for NeuraSoftmaxLayer {

src/lib.rs

@@ -1,6 +1,7 @@
 #![feature(generic_arg_infer)]
 // #![feature(generic_const_exprs)]
 #![feature(associated_type_defaults)]
+#![feature(arc_unwrap_or_clone)]
 
 pub mod algebra;
 pub mod derivable;

src/network/mod.rs

@@ -2,6 +2,7 @@ use crate::{
     algebra::NeuraVectorSpace, gradient_solver::NeuraGradientSolverBase, layer::NeuraLayer,
 };
 
+pub mod residual;
 pub mod sequential;
 
 // TODO: extract regularize from this, so that we can drop the trait constraints on NeuraSequential's impl

src/network/residual/axis.rs

@@ -0,0 +1,63 @@
+use std::borrow::Borrow;
+
+use nalgebra::{Const, DVector, Dyn, VecStorage};
+
+use crate::prelude::NeuraShape;
+
+#[derive(Clone, Copy, Debug)]
+pub struct NeuraAxisAppend;
+
+#[derive(Clone, Copy, Debug)]
+pub enum NeuraAxisErr {
+    NoInput,
+    ConflictingShape(NeuraShape, NeuraShape),
+}
+
+pub trait NeuraCombineInputs<T> {
+    type Combined;
+
+    fn combine(&self, inputs: Vec<impl Borrow<T>>) -> Self::Combined;
+
+    // TODO:
+    // fn shape(&self, input_shapes: Vec<NeuraShape>) -> NeuraShape;
+}
+
+impl<F: Clone> NeuraCombineInputs<DVector<F>> for NeuraAxisAppend {
+    type Combined = DVector<F>;
+
+    fn combine(&self, inputs: Vec<impl Borrow<DVector<F>>>) -> Self::Combined {
+        assert!(inputs.len() > 0);
+        let mut res = Vec::with_capacity(inputs.iter().map(|vec| vec.borrow().len()).sum());
+
+        for input in inputs {
+            for x in input.borrow().iter() {
+                res.push(x.clone());
+            }
+        }
+
+        DVector::from_data(VecStorage::new(Dyn(res.len()), Const as Const<1>, res))
+    }
+}
+
+impl NeuraCombineInputs<NeuraShape> for NeuraAxisAppend {
+    type Combined = Result<NeuraShape, NeuraAxisErr>;
+
+    fn combine(&self, inputs: Vec<impl Borrow<NeuraShape>>) -> Self::Combined {
+        let mut inputs = inputs.into_iter().map(|x| *x.borrow());
+        if let Some(mut res) = inputs.next() {
+            for operand in inputs {
+                match (res, operand) {
+                    (NeuraShape::Vector(x), NeuraShape::Vector(y)) => {
+                        res = NeuraShape::Vector(x + y);
+                    }
+                    (x, y) => {
+                        return Err(NeuraAxisErr::ConflictingShape(x, y));
+                    }
+                }
+            }
+            Ok(res)
+        } else {
+            Err(NeuraAxisErr::NoInput)
+        }
+    }
+}

src/network/residual/construct.rs

@@ -0,0 +1,153 @@
+use super::*;
+
+pub trait NeuraResidualConstruct {
+    type Constructed;
+    type Err;
+
+    fn construct_residual(
+        self,
+        input: NeuraResidualInput<NeuraShape>,
+    ) -> Result<Self::Constructed, Self::Err>;
+}
+
+#[derive(Clone, Debug)]
+pub enum NeuraResidualConstructErr<LayerErr, ChildErr> {
+    LayerErr(LayerErr),
+    ChildErr(ChildErr),
+    OOBConnection(usize),
+    AxisErr(NeuraAxisErr),
+}
+
+use NeuraResidualConstructErr::*;
+
+impl<Layer: NeuraPartialLayer, Axis> NeuraResidualConstruct for NeuraResidualNode<Layer, (), Axis>
+where
+    Axis: NeuraCombineInputs<NeuraShape, Combined = Result<NeuraShape, NeuraAxisErr>>,
+{
+    type Constructed = NeuraResidualNode<Layer::Constructed, (), Axis>;
+    type Err = NeuraResidualConstructErr<Layer::Err, ()>;
+
+    fn construct_residual(
+        self,
+        input: NeuraResidualInput<NeuraShape>,
+    ) -> Result<Self::Constructed, Self::Err> {
+        let (layer_input_shape, _rest) = input.shift();
+        let layer_input_shape = self
+            .axis
+            .combine(layer_input_shape)
+            .map_err(|e| AxisErr(e))?;
+
+        let layer = self
+            .layer
+            .construct(layer_input_shape)
+            .map_err(|e| LayerErr(e))?;
+        let layer_shape = layer.output_shape();
+
+        if let Some(oob_offset) = self.offsets.iter().copied().find(|o| *o > 0) {
+            return Err(OOBConnection(oob_offset));
+        }
+        // TODO: check rest for non-zero columns
+
+        Ok(NeuraResidualNode {
+            layer,
+            child_network: (),
+            offsets: self.offsets,
+            axis: self.axis,
+            output_shape: Some(layer_shape),
+        })
+    }
+}
+
+impl<Layer: NeuraPartialLayer, ChildNetwork: NeuraResidualConstruct, Axis> NeuraResidualConstruct
+    for NeuraResidualNode<Layer, ChildNetwork, Axis>
+where
+    Axis: NeuraCombineInputs<NeuraShape, Combined = Result<NeuraShape, NeuraAxisErr>>,
+{
+    type Constructed = NeuraResidualNode<Layer::Constructed, ChildNetwork::Constructed, Axis>;
+    type Err = NeuraResidualConstructErr<Layer::Err, ChildNetwork::Err>;
+
+    fn construct_residual(
+        self,
+        input: NeuraResidualInput<NeuraShape>,
+    ) -> Result<Self::Constructed, Self::Err> {
+        let (layer_input_shape, mut rest) = input.shift();
+        let layer_input_shape = self
+            .axis
+            .combine(layer_input_shape)
+            .map_err(|e| AxisErr(e))?;
+
+        let layer = self
+            .layer
+            .construct(layer_input_shape)
+            .map_err(|e| LayerErr(e))?;
+        let layer_shape = Rc::new(layer.output_shape());
+
+        for &offset in &self.offsets {
+            rest.push(offset, Rc::clone(&layer_shape));
+        }
+        let layer_shape = *layer_shape;
+
+        let child_network = self
+            .child_network
+            .construct_residual(rest)
+            .map_err(|e| ChildErr(e))?;
+
+        Ok(NeuraResidualNode {
+            layer,
+            child_network,
+            offsets: self.offsets,
+            axis: self.axis,
+            output_shape: Some(layer_shape),
+        })
+    }
+}
+
+impl<Layer, Axis> NeuraShapedLayer for NeuraResidualNode<Layer, (), Axis> {
+    #[inline(always)]
+    fn output_shape(&self) -> NeuraShape {
+        self.output_shape.unwrap()
+    }
+}
+
+impl<Layer, ChildNetwork: NeuraShapedLayer, Axis> NeuraShapedLayer
+    for NeuraResidualNode<Layer, ChildNetwork, Axis>
+{
+    #[inline(always)]
+    fn output_shape(&self) -> NeuraShape {
+        self.child_network.output_shape()
+    }
+}
+
+impl<Layers: NeuraShapedLayer> NeuraShapedLayer for NeuraResidual<Layers> {
+    #[inline(always)]
+    fn output_shape(&self) -> NeuraShape {
+        self.layers.output_shape()
+    }
+}
+
+impl<Layers: NeuraResidualConstruct> NeuraPartialLayer for NeuraResidual<Layers>
+where
+    // Should always be satisfied:
+    Layers::Constructed: NeuraShapedLayer,
+{
+    type Constructed = NeuraResidual<Layers::Constructed>;
+    type Err = Layers::Err;
+
+    fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
+        let input_shape = Rc::new(input_shape);
+        let mut inputs = NeuraResidualInput::new();
+
+        for &offset in &self.initial_offsets {
+            inputs.push(offset, Rc::clone(&input_shape));
+        }
+
+        drop(input_shape);
+
+        let layers = self.layers.construct_residual(inputs)?;
+
+        Ok(NeuraResidual {
+            layers,
+            initial_offsets: self.initial_offsets,
+        })
+    }
+}

src/network/residual/input.rs

@@ -0,0 +1,38 @@
+use std::rc::Rc;
+
+#[derive(Clone, Debug)]
+pub struct NeuraResidualInput<T> {
+    // TODO: try to remove this Rc
+    slots: Vec<Vec<Rc<T>>>,
+}
+
+impl<T> NeuraResidualInput<T> {
+    pub fn new() -> Self {
+        Self { slots: Vec::new() }
+    }
+
+    pub fn push(&mut self, offset: usize, data: Rc<T>) {
+        while self.slots.len() <= offset {
+            self.slots.push(Vec::new());
+        }
+        self.slots[offset].push(data);
+    }
+
+    pub fn shift(&self) -> (Vec<Rc<T>>, NeuraResidualInput<T>) {
+        let res = self.slots.get(0).cloned().unwrap_or(vec![]);
+        let new_input = Self {
+            slots: self.slots.iter().skip(1).cloned().collect(),
+        };
+
+        (res, new_input)
+    }
+
+    /// Returns the first input item of the first slot
+    pub fn get_first(self) -> Option<Rc<T>> {
+        // TODO: return None if the first slot is bigger than 1 or if there are multiple non-empty slots
+        self.slots
+            .into_iter()
+            .next()
+            .and_then(|first_slot| first_slot.into_iter().next())
+    }
+}

src/network/residual/mod.rs

@@ -0,0 +1,259 @@
+use std::rc::Rc;
+
+use nalgebra::{DVector, Scalar};
+use num::Float;
+
+use crate::layer::*;
+
+mod input;
+pub use input::*;
+
+mod axis;
+pub use axis::*;
+
+mod construct;
+pub use construct::NeuraResidualConstructErr;
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct NeuraResidual<Layers> {
+    /// Instance of NeuraResidualNode
+    layers: Layers,
+
+    /// Array of which layers to send the input to, defaults to `vec![0]`
+    initial_offsets: Vec<usize>,
+}
+
+impl<Layers> NeuraResidual<Layers> {
+    pub fn new(layers: Layers) -> Self {
+        Self {
+            layers,
+            initial_offsets: vec![0],
+        }
+    }
+
+    pub fn offset(mut self, offset: usize) -> Self {
+        self.initial_offsets.push(offset);
+        self
+    }
+
+    pub fn offsets(mut self, offsets: Vec<usize>) -> Self {
+        self.initial_offsets = offsets;
+        self
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct NeuraResidualNode<Layer, ChildNetwork, Axis> {
+    pub layer: Layer,
+    pub child_network: ChildNetwork,
+
+    /// Array of relative layers indices to send the offset of this layer to,
+    /// defaults to `vec![0]`.
+    offsets: Vec<usize>,
+
+    pub axis: Axis,
+
+    output_shape: Option<NeuraShape>,
+}
+
+impl<Layer, ChildNetwork> NeuraResidualNode<Layer, ChildNetwork, NeuraAxisAppend> {
+    pub fn new(layer: Layer, child_network: ChildNetwork) -> Self {
+        Self {
+            layer,
+            child_network,
+            offsets: vec![0],
+            axis: NeuraAxisAppend,
+            output_shape: None,
+        }
+    }
+}
+
+impl<Layer, ChildNetwork, Axis> NeuraResidualNode<Layer, ChildNetwork, Axis> {
+    pub fn offsets(mut self, offsets: Vec<usize>) -> Self {
+        self.offsets = offsets;
+        self
+    }
+
+    pub fn offset(mut self, offset: usize) -> Self {
+        self.offsets.push(offset);
+        self
+    }
+
+    pub fn axis<Axis2>(self, axis: Axis2) -> NeuraResidualNode<Layer, ChildNetwork, Axis2> {
+        NeuraResidualNode {
+            layer: self.layer,
+            child_network: self.child_network,
+            offsets: self.offsets,
+            axis,
+            // Drop the knowledge of output_shape
+            output_shape: None,
+        }
+    }
+}
+
+impl<F: Float + Scalar, Layer, ChildNetwork, Axis> NeuraLayer<NeuraResidualInput<DVector<F>>>
+    for NeuraResidualNode<Layer, ChildNetwork, Axis>
+where
+    Axis: NeuraCombineInputs<DVector<F>>,
+    Layer: NeuraLayer<Axis::Combined, Output = DVector<F>>,
+    ChildNetwork: NeuraLayer<NeuraResidualInput<DVector<F>>>,
+{
+    type Output = <ChildNetwork as NeuraLayer<NeuraResidualInput<DVector<F>>>>::Output;
+
+    fn eval(&self, input: &NeuraResidualInput<DVector<F>>) -> Self::Output {
+        let (inputs, mut rest) = input.shift();
+
+        let layer_input = self.axis.combine(inputs);
+        let layer_output = Rc::new(self.layer.eval(&layer_input));
+
+        for &offset in &self.offsets {
+            rest.push(offset, Rc::clone(&layer_output));
+        }
+
+        self.child_network.eval(&rest)
+    }
+}
+
+impl<F: Clone, Output: Clone, Layers> NeuraLayer<DVector<F>> for NeuraResidual<Layers>
+where
+    Layers: NeuraLayer<NeuraResidualInput<DVector<F>>, Output = NeuraResidualInput<Output>>,
+{
+    type Output = Output;
+
+    fn eval(&self, input: &DVector<F>) -> Self::Output {
+        let input: Rc<DVector<F>> = Rc::new((*input).clone());
+        let mut inputs = NeuraResidualInput::new();
+
+        for &offset in &self.initial_offsets {
+            inputs.push(offset, Rc::clone(&input));
+        }
+
+        drop(input);
+
+        let output = self.layers.eval(&inputs);
+
+        let result = output.get_first()
+            .expect("Invalid NeuraResidual state: network returned no data, did you forget to link the last layer?")
+            .into();
+
+        Rc::unwrap_or_clone(result)
+    }
+}
+
+#[macro_export]
+macro_rules! neura_residual {
+    [ "__combine_layers", ] => {
+        ()
+    };
+
+    [ "__combine_layers",
+        $layer:expr $(, $axis:expr)? $( => $( $offset:expr ),* )?
+        $(; $( $rest_layer:expr $(, $rest_axis:expr)? $( => $( $rest_offset:expr ),* )? );*)?
+    ] => {{
+        let layer = $crate::network::residual::NeuraResidualNode::new($layer,
+            neura_residual![
+                "__combine_layers",
+                $($( $rest_layer $(, $rest_axis)? $( => $( $rest_offset ),* )? );*)?
+            ]
+        );
+
+        $(
+            let layer = layer.axis($axis);
+        )?
+
+        $(
+            let layer = layer.offsets(vec![$($offset),*]);
+        )?
+
+        layer
+    }};
+
+    [
+        $( <= $( $initial_offset:expr ),* ;)?
+        $( $layer:expr $(, $axis:expr)? $( => $( $offset:expr ),* $(,)? )? );*
+        $(;)?
+    ] => {{
+        let res = $crate::network::residual::NeuraResidual::new(
+            neura_residual![ "__combine_layers", $( $layer $(, $axis)? $( => $( $offset ),* )? );* ]
+        );
+
+        $(
+            let res = res.offsets(vec![$($initial_offset),*]);
+        )?
+
+        res
+    }};
+}
+
+#[cfg(test)]
+mod test {
+    use nalgebra::dvector;
+
+    use crate::neura_layer;
+
+    use super::*;
+
+    #[test]
+    fn test_resnet_eval() {
+        let network = NeuraResidual::new(
+            NeuraResidualNode::new(
+                neura_layer!("dense", 4)
+                    .construct(NeuraShape::Vector(2))
+                    .unwrap(),
+                NeuraResidualNode::new(
+                    neura_layer!("dense", 3)
+                        .construct(NeuraShape::Vector(4))
+                        .unwrap(),
+                    NeuraResidualNode::new(
+                        neura_layer!("dense", 6)
+                            .construct(NeuraShape::Vector(7))
+                            .unwrap(),
+                        (),
+                    ),
+                ),
+            )
+            .offset(1),
+        );
+
+        network.eval(&dvector![0.2, 0.4]);
+    }
+
+    #[test]
+    fn test_resnet_macro() {
+        let network = neura_residual![
+            <= 0, 2;
+            neura_layer!("dense", 5) => 0, 1;
+            neura_layer!("dense", 5);
+            neura_layer!("dense", 3)
+        ];
+
+        println!("{:#?}", network);
+
+        assert_eq!(network.initial_offsets, vec![0, 2]);
+        assert_eq!(network.layers.offsets, vec![0, 1]);
+        assert_eq!(network.layers.child_network.offsets, vec![0]);
+        assert_eq!(network.layers.child_network.child_network.child_network, ());
+
+        let network = neura_residual![
+            neura_layer!("dense", 4) => 0;
+        ];
+
+        assert_eq!(network.initial_offsets, vec![0]);
+    }
+
+    #[test]
+    fn test_resnet_partial() {
+        let network = neura_residual![
+            <= 0, 1;
+            neura_layer!("dense", 2) => 0, 1;
+            neura_layer!("dense", 4);
+            neura_layer!("dense", 8)
+        ]
+        .construct(NeuraShape::Vector(1))
+        .unwrap();
+
+        assert_eq!(network.output_shape(), NeuraShape::Vector(8));
+
+        network.eval(&dvector![0.0]);
+    }
+}

src/network/sequential/construct.rs

@@ -1,3 +1,5 @@
+use crate::layer::NeuraShapedLayer;
+
 use super::*;
 
 pub trait NeuraSequentialConstruct {
@@ -40,7 +42,7 @@ impl<Layer: NeuraPartialLayer, ChildNetwork: NeuraSequentialConstruct> NeuraSequ
         // TODO: ensure that this operation (and all recursive operations) are directly allocated on the heap
         let child_network = self
             .child_network
-            .construct(Layer::output_shape(&layer))
+            .construct(layer.output_shape())
             .map_err(|e| NeuraSequentialConstructErr::Child(e))?;
         let child_network = Box::new(child_network);
 
@@ -50,3 +52,19 @@ impl<Layer: NeuraPartialLayer, ChildNetwork: NeuraSequentialConstruct> NeuraSequ
         })
     }
 }
+
+impl<Layer: NeuraShapedLayer> NeuraShapedLayer for NeuraSequential<Layer, ()> {
+    #[inline(always)]
+    fn output_shape(&self) -> NeuraShape {
+        self.layer.output_shape()
+    }
+}
+
+impl<Layer, ChildNetwork: NeuraShapedLayer> NeuraShapedLayer
+    for NeuraSequential<Layer, ChildNetwork>
+{
+    #[inline(always)]
+    fn output_shape(&self) -> NeuraShape {
+        self.child_network.output_shape()
+    }
+}