onnx-export-rs 0.1.0

Export canonical Rust machine-learning models to ONNX
Documentation
use ndarray::{Array1, Array2};

use crate::{Error, Result};

/// Fitted categorical Naive Bayes log-probability tables.
#[derive(Clone, Debug, PartialEq)]
pub struct CategoricalNaiveBayes {
    /// Per-feature tables shaped `[categories, classes]`.
    pub feature_log_probabilities: Vec<Array2<f64>>,
    /// Log prior for every class.
    pub log_priors: Array1<f64>,
    /// External integer labels in class-column order.
    pub class_labels: Vec<i64>,
}

impl CategoricalNaiveBayes {
    /// Creates validated categorical Naive Bayes state.
    pub fn new(
        feature_log_probabilities: Vec<Array2<f64>>,
        log_priors: Array1<f64>,
        class_labels: Vec<i64>,
    ) -> Result<Self> {
        let classes = class_labels.len();
        if classes == 0
            || log_priors.len() != classes
            || feature_log_probabilities.is_empty()
            || feature_log_probabilities.iter().any(|table| {
                table.nrows() == 0
                    || table.ncols() != classes
                    || table
                        .iter()
                        .any(|value| value.is_nan() || *value == f64::INFINITY)
            })
            || log_priors
                .iter()
                .any(|value| value.is_nan() || *value == f64::INFINITY)
        {
            return Err(Error::InvalidModel(
                "invalid categorical Naive Bayes state".into(),
            ));
        }
        Ok(Self {
            feature_log_probabilities,
            log_priors,
            class_labels,
        })
    }
}