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}