onnx_graph 0.1.2

use std::any::Any;

use crate::{
    nodes::{node::Node, onnx_operation_trait::FromOnnxOperation, unique_ids::UniqueId},
    tensor_map::TensorMap,
    typed_array::TypedArray,
};
use anyhow::Result;
use onnx_extractor::OnnxOperation;

#[derive(Default)]
pub struct SoftMaxNode<T: Default> {
    input: String,

    o: String,

    unique_id: UniqueId,

    axis: i64,

    next_node: Option<Vec<Box<dyn Node<T>>>>,
}

impl<T: Default> FromOnnxOperation for SoftMaxNode<T> {
    fn from_onnx_operation(elem: &OnnxOperation) -> Result<Self> {
        let attrs = &elem.attributes;
        let mut softmax = Self {
            input: String::new(),
            o: String::new(),
            axis: match attrs.get("axis") {
                Some(av) => av.as_int().unwrap(),
                None => 0,
            },
            unique_id: UniqueId::Softmax,
            next_node: None,
        };
        softmax.add_input_strings(elem.inputs[0].clone());
        softmax.add_output_strings(elem.outputs[0].clone());
        Ok(softmax)
    }
}

impl<T: Default> SoftMaxNode<T> {
    pub fn new(axis: i64) -> Self {
        Self {
            input: String::new(),

            o: String::new(),

            unique_id: UniqueId::Softmax,

            axis,
            next_node: None,
        }
    }

    pub fn add_input_strings(&mut self, input: String) {
        self.input = input;
    }

    pub fn add_output_strings(&mut self, o: String) {
        self.o = o;
    }
}

impl<T: Default + 'static> Node<T> for SoftMaxNode<T> {
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn get_unique_id(&self) -> UniqueId {
        self.unique_id
    }
    fn get_unique_id_mut(&mut self) -> UniqueId {
        self.unique_id
    }

    fn get_next(&self) -> Option<&Vec<Box<dyn Node<T>>>> {
        self.next_node.as_ref()
    }

    fn execute(&self, omap: &mut TensorMap) {
        let [x, o] = omap.get_disjoint_mut([&self.input, &self.o]);
        let x = &*x.unwrap();

        match o {
            Some(result) => {
                x.softmax(self.axis, result).unwrap();
            }
            None => panic!("SoftMaxNode: missing input {}", self.input),
        }
    }

    fn output_names(&self) -> Vec<String> {
        vec![self.o.clone()]
    }

    fn input_names(&self) -> Vec<String> {
        vec![self.input.clone()]
    }

    fn take_next(&mut self) -> Option<Vec<Box<dyn Node<T>>>> {
        self.next_node.take()
    }
    fn get_next_mut(&mut self) -> Option<&mut Vec<Box<dyn Node<T>>>> {
        self.next_node.as_mut()
    }

    fn set_next(&mut self, next: Option<Vec<Box<dyn Node<T>>>>) {
        self.next_node = next;
    }

    fn print(&self) {
        if let Some(list) = &self.next_node {
            print!("{}-", list.len());
        }
        println!("soft_max-{},{}", self.input, self.o);
        if let Some(next) = &self.next_node {
            next.iter().for_each(|v| v.print());
        }
    }

    fn determine_output_shape(&mut self, omap: &mut TensorMap) {
        let [x, o] = omap.get_disjoint_mut([&self.input, &self.o]);
        let x = x.map(|arr| &*arr);

        if let (Some(x), Some(o)) = (x, o)
            && let Some(in_shape) = x.shape()
        {
            *o = TypedArray::empty_with_others_type(x, in_shape);
        }

        if let Some(list) = &mut self.next_node {
            for next in list {
                next.determine_output_shape(omap);
            }
        }
    }
}

macro_rules! softmax_variant {
    ($variant:ident, $axis:expr, $a:expr, $o:expr, $T:ty) => {{
        use ndarray::ArrayD;
        use ndarray::Axis;
        use ndarray::IxDyn;

        let ndim = $a.ndim() as i64;
        let axis = if $axis < 0 {
            (ndim + $axis) as usize
        } else {
            $axis as usize
        };

        let needs_alloc = match &*$o {
            TypedArray::$variant(out) => out.shape() != $a.shape(),
            _ => true,
        };
        if needs_alloc {
            *$o = TypedArray::$variant(ArrayD::zeros(IxDyn($a.shape())));
        }

        if let TypedArray::$variant(out) = $o {
            let dst = out.as_slice_memory_order_mut().unwrap();
            let src = $a.as_slice_memory_order().unwrap();
            dst.copy_from_slice(src);

            for mut lane in out.lanes_mut(Axis(axis)) {
                let max = lane.iter().copied().fold(<$T>::NEG_INFINITY, <$T>::max);
                lane.mapv_inplace(|x| (x - max).exp());
                let sum: $T = lane.iter().sum();
                lane.mapv_inplace(|x| x / sum);
            }
        }
    }};
}

impl TypedArray {
    pub fn softmax(&self, axis: i64, o: &mut TypedArray) -> anyhow::Result<()> {
        match self {
            TypedArray::Float(a) => softmax_variant!(Float, axis, a, o, f32),
            TypedArray::Double(a) => softmax_variant!(Double, axis, a, o, f64),
            _ => return Err(anyhow::anyhow!("softmax only supported for F32/F64")),
        }
        Ok(())
    }
}