native_neural_network_std 0.2.1

use native_neural_network::rnn_api::RnnApiError;

pub struct ModelStd {
    pub layers: Vec<crate::std::layers_std::LayerSpec>,
    pub weights: Vec<f32>,
    pub biases: Vec<f32>,
    pub input_size: usize,
    pub output_size: usize,
    pub infer_scratch_len: usize,
}

impl ModelStd {
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, RnnApiError> {
        let counts = crate::std::rnn_std::rnn_required_from_bytes_v1_std(bytes)?;
        if counts.layers == 0 {
            return Err(RnnApiError::BadBytes);
        }

        let (layers, weights, biases) = crate::std::rnn_std::rnn_unpack_v1_std(bytes)?;

        let input_size = layers
            .first()
            .map(|l| l.input_size())
            .ok_or(RnnApiError::BadBytes)?;
        let output_size = layers
            .last()
            .map(|l| l.output_size())
            .ok_or(RnnApiError::BadBytes)?;

        let infer_scratch_len =
            crate::std::rnn_std::rnn_required_infer_scratch_from_specs_std(&layers)?;

        Ok(Self {
            layers,
            weights,
            biases,
            input_size,
            output_size,
            infer_scratch_len,
        })
    }

    pub fn run_single(&self, input: &[f32]) -> Result<Vec<f32>, RnnApiError> {
        if input.len() != self.input_size {
            return Err(RnnApiError::InvalidTopology);
        }
        let mut infer_scratch = vec![0f32; self.infer_scratch_len];
        let mut output = vec![0f32; self.output_size];
        let plan_std = crate::std::layers_std::LayerPlanStd::new(
            self.layers.clone(),
            self.weights.clone(),
            self.biases.clone(),
        );
        crate::std::engine_std::forward_plan(&plan_std, input, &mut output, &mut infer_scratch)
            .map_err(|_| RnnApiError::Model)?;
        Ok(output)
    }

    pub fn run_batch(&self, input: &[f32], batch_size: usize) -> Result<Vec<f32>, RnnApiError> {
        if batch_size == 0 {
            return Err(RnnApiError::InvalidTopology);
        }
        let expected_in = self
            .input_size
            .checked_mul(batch_size)
            .ok_or(RnnApiError::InvalidTopology)?;
        let expected_out = self
            .output_size
            .checked_mul(batch_size)
            .ok_or(RnnApiError::InvalidTopology)?;
        if input.len() != expected_in {
            return Err(RnnApiError::InvalidTopology);
        }

        let mut infer_scratch = vec![0f32; self.infer_scratch_len];
        let mut output = vec![0f32; expected_out];
        let plan_std = crate::std::layers_std::LayerPlanStd::new(
            self.layers.clone(),
            self.weights.clone(),
            self.biases.clone(),
        );
        for row in 0..batch_size {
            let in_off = row * self.input_size;
            let out_off = row * self.output_size;
            crate::std::engine_std::forward_plan(
                &plan_std,
                &input[in_off..in_off + self.input_size],
                &mut output[out_off..out_off + self.output_size],
                &mut infer_scratch,
            )
            .map_err(|_| RnnApiError::Model)?;
        }

        Ok(output)
    }

    pub fn info(&self) -> (usize, usize, usize) {
        (self.layers.len(), self.weights.len(), self.biases.len())
    }
}