concision_neural/traits/
predict.rs1use cnc::Forward;
6
7pub trait Predict<Rhs> {
16    type Output;
17
18    private!();
19
20    fn predict(&self, input: &Rhs) -> crate::ModelResult<Self::Output>;
21}
22
23pub trait PredictWithConfidence<Rhs>: Predict<Rhs> {
26    type Confidence;
27
28    fn predict_with_confidence(
29        &self,
30        input: &Rhs,
31    ) -> crate::ModelResult<(Self::Output, Self::Confidence)>;
32}
33
34use ndarray::{Array, Dimension, ScalarOperand};
39use num_traits::{Float, FromPrimitive};
40
41impl<M, U, V> Predict<U> for M
42where
43    M: Forward<U, Output = V>,
44{
45    type Output = V;
46
47    seal!();
48
49    fn predict(&self, input: &U) -> crate::ModelResult<Self::Output> {
50        self.forward(input).map_err(core::convert::Into::into)
51    }
52}
53
54impl<M, U, A, D> PredictWithConfidence<U> for M
55where
56    A: Float + FromPrimitive + ScalarOperand,
57    D: Dimension,
58    Self: Predict<U, Output = Array<A, D>>,
59{
60    type Confidence = A;
61
62    fn predict_with_confidence(
63        &self,
64        input: &U,
65    ) -> Result<(Self::Output, Self::Confidence), crate::ModelError> {
66        let prediction = Predict::predict(self, input)?;
68        let shape = prediction.shape();
69        let batch_size = shape[0];
72
73        let mut variance_sum = A::zero();
74
75        for sample in prediction.rows() {
76            let variance = sample.var(A::one());
78            variance_sum = variance_sum + variance;
79        }
80
81        let avg_variance = variance_sum / A::from_usize(batch_size).unwrap();
83        let confidence = (A::one() + avg_variance).recip();
85
86        Ok((prediction, confidence))
87    }
88}