eredu_runtime/
decision.rs

1//! Backend-neutral sequential prediction decisions and layered traversal handoff.
2
3use std::marker::PhantomData;
4
5use eredu_nn::{NeuralBackend, Tensor};
6
7use crate::{
8    layered::{LayeredTraversalHook, LayeredTraversalPoint, LayeredUnitAction},
9    Sampler, SamplingBackend, TokenDomain,
10};
11
12/// How a complete sequential prediction plan obtains its tokens.
13#[derive(Debug, Clone, Copy, Eq, PartialEq)]
14pub enum SequentialDecisionMode {
15    /// Every prediction is supplied by the caller.
16    TeacherForced,
17    /// Every prediction is selected from backend-native logits.
18    Autoregressive,
19    /// Forced and sampler-selected predictions are interleaved.
20    PartiallyForced,
21}
22
23/// One prediction's forcing directive.
24#[derive(Debug, Clone, Eq, PartialEq)]
25pub enum PredictionDirective<T> {
26    /// Select a token from this prediction's logits.
27    Sample,
28    /// Supply this backend-native token without invoking the sampler.
29    Force(T),
30}
31
32/// Validated forcing and diagnostic policy for one ordered prediction chain.
33#[derive(Debug, Clone, Eq, PartialEq)]
34pub struct SequentialDecisionPlan<T> {
35    directives: Vec<PredictionDirective<T>>,
36    retain_diagnostics: bool,
37    allow_fully_forced_tail_skip: bool,
38}
39
40impl<T> SequentialDecisionPlan<T> {
41    /// Creates a non-empty ordered prediction plan.
42    pub fn new(
43        directives: impl IntoIterator<Item = PredictionDirective<T>>,
44        retain_diagnostics: bool,
45        allow_fully_forced_tail_skip: bool,
46    ) -> Result<Self, SequentialDecisionPlanError> {
47        let directives = directives.into_iter().collect::<Vec<_>>();
48        if directives.is_empty() {
49            return Err(SequentialDecisionPlanError::EmptyPlan);
50        }
51        Ok(Self {
52            directives,
53            retain_diagnostics,
54            allow_fully_forced_tail_skip,
55        })
56    }
57
58    /// Returns the number of ordered target and predictor decisions.
59    pub fn len(&self) -> usize {
60        self.directives.len()
61    }
62
63    /// Returns whether this plan contains no decisions.
64    pub fn is_empty(&self) -> bool {
65        self.directives.is_empty()
66    }
67
68    /// Returns the aggregate forcing mode.
69    pub fn mode(&self) -> SequentialDecisionMode {
70        let forced = self
71            .directives
72            .iter()
73            .filter(|directive| matches!(directive, PredictionDirective::Force(_)))
74            .count();
75        match forced {
76            0 => SequentialDecisionMode::Autoregressive,
77            count if count == self.directives.len() => SequentialDecisionMode::TeacherForced,
78            _ => SequentialDecisionMode::PartiallyForced,
79        }
80    }
81
82    /// Returns a portable per-prediction forcing mask in decision order.
83    pub fn forcing_mask(&self) -> impl ExactSizeIterator<Item = bool> + '_ {
84        self.directives
85            .iter()
86            .map(|directive| matches!(directive, PredictionDirective::Force(_)))
87    }
88
89    /// Returns whether diagnostic logits are retained for every executed decision.
90    pub const fn retains_diagnostics(&self) -> bool {
91        self.retain_diagnostics
92    }
93
94    /// Returns whether a proven fully forced tail may omit model calls.
95    pub const fn allows_fully_forced_tail_skip(&self) -> bool {
96        self.allow_fully_forced_tail_skip
97    }
98}
99
100/// Origin of one resolved sequential token.
101#[derive(Debug, Clone, Copy, Eq, PartialEq)]
102pub enum SequentialDecisionSource {
103    /// The caller forced the token after logits were computed.
104    Forced,
105    /// A backend-native sampler selected the token.
106    Sampled,
107    /// The caller forced the token in a tail whose model calls were omitted.
108    ForcedTailSkipped,
109}
110
111/// One resolved token in prediction order.
112#[derive(Debug, Clone, Eq, PartialEq)]
113pub struct SequentialDecision<T> {
114    prediction: usize,
115    source: SequentialDecisionSource,
116    token: T,
117}
118
119impl<T> SequentialDecision<T> {
120    /// Returns the zero-based prediction ordinal.
121    pub const fn prediction(&self) -> usize {
122        self.prediction
123    }
124
125    /// Returns whether forcing or sampling produced this token.
126    pub const fn source(&self) -> SequentialDecisionSource {
127        self.source
128    }
129
130    /// Borrows the backend-native selected token.
131    pub const fn token(&self) -> &T {
132        &self.token
133    }
134}
135
136/// Diagnostic logits retained at one executed decision boundary.
137#[derive(Debug, Clone, Eq, PartialEq)]
138pub struct SequentialDecisionDiagnostic<L> {
139    prediction: usize,
140    logits: L,
141}
142
143impl<L> SequentialDecisionDiagnostic<L> {
144    /// Returns the zero-based prediction ordinal.
145    pub const fn prediction(&self) -> usize {
146        self.prediction
147    }
148
149    /// Borrows the backend-native logits without host materialization.
150    pub const fn logits(&self) -> &L {
151        &self.logits
152    }
153}
154
155/// Validated result of checking whether the unexecuted group tail is skippable.
156#[derive(Debug, Clone, Copy, Eq, PartialEq)]
157pub enum FullyForcedTailDecision {
158    /// At least one remaining unit must execute.
159    Execute,
160    /// Every remaining unit corresponds exactly to one forced decision.
161    Skip {
162        /// Number of forced predictions proven safe to omit.
163        predictions: usize,
164    },
165}
166
167/// Statically dispatched sequential token resolver.
168///
169/// The driver owns one sampler per prediction and one optional backend random
170/// state. Forced tokens and diagnostic logits remain backend-native values.
171pub struct SequentialDecisionDriver<B, S>
172where
173    B: SamplingBackend,
174    S: Sampler<B>,
175{
176    plan: SequentialDecisionPlan<B::Token>,
177    samplers: Vec<S>,
178    temperatures: Vec<f32>,
179    random: Option<B::RandomState>,
180    decisions: Vec<SequentialDecision<B::Token>>,
181    diagnostics: Vec<SequentialDecisionDiagnostic<B::Logits>>,
182}
183
184/// Sampler instances and optional backend randomness advanced by one decision pass.
185pub type SequentialSamplingState<S, R> = (Vec<S>, Option<R>);
186
187impl<B, S> SequentialDecisionDriver<B, S>
188where
189    B: SamplingBackend,
190    S: Sampler<B>,
191{
192    /// Creates a driver with exactly one sampler and temperature per prediction.
193    pub fn new(
194        plan: SequentialDecisionPlan<B::Token>,
195        samplers: Vec<S>,
196        temperatures: Vec<f32>,
197        random: Option<B::RandomState>,
198    ) -> Result<Self, SequentialDecisionPlanError> {
199        if samplers.len() != plan.len() {
200            return Err(SequentialDecisionPlanError::SamplerCountMismatch {
201                predictions: plan.len(),
202                samplers: samplers.len(),
203            });
204        }
205        if temperatures.len() != plan.len() {
206            return Err(SequentialDecisionPlanError::TemperatureCountMismatch {
207                predictions: plan.len(),
208                temperatures: temperatures.len(),
209            });
210        }
211        if let Some((prediction, temperature)) = temperatures
212            .iter()
213            .copied()
214            .enumerate()
215            .find(|(_, temperature)| !temperature.is_finite() || *temperature < 0.0)
216        {
217            return Err(SequentialDecisionPlanError::InvalidTemperature {
218                prediction,
219                bits: temperature.to_bits(),
220            });
221        }
222        Ok(Self {
223            plan,
224            samplers,
225            temperatures,
226            random,
227            decisions: Vec::new(),
228            diagnostics: Vec::new(),
229        })
230    }
231
232    /// Borrows the validated decision plan.
233    pub const fn plan(&self) -> &SequentialDecisionPlan<B::Token> {
234        &self.plan
235    }
236
237    /// Returns the next prediction ordinal expected at a traversal boundary.
238    pub fn next_prediction(&self) -> usize {
239        self.decisions.len()
240    }
241
242    /// Returns resolved decisions in canonical order.
243    pub fn decisions(&self) -> &[SequentialDecision<B::Token>] {
244        &self.decisions
245    }
246
247    /// Returns retained diagnostic logits in canonical order.
248    pub fn diagnostics(&self) -> &[SequentialDecisionDiagnostic<B::Logits>] {
249        &self.diagnostics
250    }
251
252    /// Borrows the backend random state after all decisions made so far.
253    pub const fn random_state(&self) -> Option<&B::RandomState> {
254        self.random.as_ref()
255    }
256
257    /// Validates whether `remaining_units` is exactly one fully forced tail.
258    pub fn fully_forced_tail_decision(
259        &self,
260        prediction: usize,
261        remaining_units: usize,
262    ) -> Result<FullyForcedTailDecision, SequentialDecisionError<B::Error>> {
263        self.require_next(prediction)?;
264        let remaining_predictions = self.plan.len().saturating_sub(prediction);
265        if !self.plan.allow_fully_forced_tail_skip
266            || self.plan.retain_diagnostics
267            || remaining_predictions != remaining_units
268            || !self.plan.directives[prediction..]
269                .iter()
270                .all(|directive| matches!(directive, PredictionDirective::Force(_)))
271        {
272            return Ok(FullyForcedTailDecision::Execute);
273        }
274        Ok(FullyForcedTailDecision::Skip {
275            predictions: remaining_predictions,
276        })
277    }
278
279    /// Returns cloned forced tokens for a tail already proven skippable.
280    pub fn forced_tail_tokens(
281        &self,
282        prediction: usize,
283        count: usize,
284        domains: impl IntoIterator<Item = TokenDomain>,
285        context: &B::Context,
286    ) -> Result<Vec<B::Token>, SequentialDecisionError<B::Error>> {
287        self.require_next(prediction)?;
288        if self.fully_forced_tail_decision(prediction, count)?
289            != (FullyForcedTailDecision::Skip { predictions: count })
290        {
291            return Err(SequentialDecisionError::InvalidTailSkip { prediction, count });
292        }
293        let domains = domains.into_iter().collect::<Vec<_>>();
294        if domains.len() != count {
295            return Err(SequentialDecisionError::TokenDomainCountMismatch {
296                prediction,
297                expected: count,
298                actual: domains.len(),
299            });
300        }
301        self.plan.directives[prediction..prediction + count]
302            .iter()
303            .zip(domains)
304            .map(|(directive, domain)| match directive {
305                PredictionDirective::Force(token) => B::validate_token(token, domain, context)
306                    .map_err(SequentialDecisionError::Backend),
307                PredictionDirective::Sample => unreachable!("tail was proven fully forced"),
308            })
309            .collect()
310    }
311
312    /// Records a proven and architecture-accepted forced tail.
313    pub fn commit_forced_tail(
314        &mut self,
315        prediction: usize,
316        tokens: Vec<B::Token>,
317    ) -> Result<(), SequentialDecisionError<B::Error>> {
318        self.require_next(prediction)?;
319        let count = tokens.len();
320        if self.fully_forced_tail_decision(prediction, count)?
321            != (FullyForcedTailDecision::Skip { predictions: count })
322        {
323            return Err(SequentialDecisionError::InvalidTailSkip { prediction, count });
324        }
325        self.decisions
326            .extend(
327                tokens
328                    .into_iter()
329                    .enumerate()
330                    .map(|(offset, token)| SequentialDecision {
331                        prediction: prediction + offset,
332                        source: SequentialDecisionSource::ForcedTailSkipped,
333                        token,
334                    }),
335            );
336        Ok(())
337    }
338
339    /// Resolves one executed prediction from backend-native logits.
340    pub fn resolve(
341        &mut self,
342        prediction: usize,
343        logits: &B::Logits,
344        domain: TokenDomain,
345        context: &B::Context,
346    ) -> Result<B::Token, SequentialDecisionError<B::Error>> {
347        self.require_next(prediction)?;
348        let (token, source) = match &self.plan.directives[prediction] {
349            PredictionDirective::Force(token) => (token.clone(), SequentialDecisionSource::Forced),
350            PredictionDirective::Sample => (
351                self.samplers[prediction]
352                    .sample(
353                        logits,
354                        self.temperatures[prediction],
355                        self.random.as_mut(),
356                        context,
357                    )
358                    .map_err(SequentialDecisionError::Backend)?,
359                SequentialDecisionSource::Sampled,
360            ),
361        };
362        let token =
363            B::validate_token(&token, domain, context).map_err(SequentialDecisionError::Backend)?;
364        if self.plan.retain_diagnostics {
365            self.diagnostics.push(SequentialDecisionDiagnostic {
366                prediction,
367                logits: logits.clone(),
368            });
369        }
370        self.decisions.push(SequentialDecision {
371            prediction,
372            source,
373            token: token.clone(),
374        });
375        Ok(token)
376    }
377
378    /// Validates that every planned prediction was resolved exactly once.
379    pub fn finish(&self) -> Result<(), SequentialDecisionError<B::Error>> {
380        if self.decisions.len() != self.plan.len() {
381            return Err(SequentialDecisionError::Incomplete {
382                resolved: self.decisions.len(),
383                predictions: self.plan.len(),
384            });
385        }
386        Ok(())
387    }
388
389    /// Finishes the existing decision sequence and returns its advanced
390    /// sampler and backend-random states for transactional publication.
391    ///
392    /// Decisions and diagnostics remain inspectable until this method consumes
393    /// the driver. A caller can therefore copy any required output metadata
394    /// before atomically adopting these state components.
395    pub fn finish_into_sampling_state(
396        self,
397    ) -> Result<SequentialSamplingState<S, B::RandomState>, SequentialDecisionError<B::Error>> {
398        self.finish()?;
399        Ok((self.samplers, self.random))
400    }
401
402    fn require_next(&self, prediction: usize) -> Result<(), SequentialDecisionError<B::Error>> {
403        if prediction != self.decisions.len() || prediction >= self.plan.len() {
404            return Err(SequentialDecisionError::OutOfOrder {
405                expected: self.decisions.len(),
406                actual: prediction,
407                predictions: self.plan.len(),
408            });
409        }
410        Ok(())
411    }
412}
413
414/// Architecture-owned conversion between layered boundaries and predictions.
415pub trait SequentialDecisionBoundary<B, C, E>
416where
417    B: SamplingBackend,
418{
419    /// Returns the prediction ordinal at this traversal point, if any.
420    fn prediction_at(&self, point: LayeredTraversalPoint, forward: &C) -> Option<usize>;
421
422    /// Produces backend-native logits for one target or predictor boundary.
423    fn logits(
424        &mut self,
425        prediction: usize,
426        point: LayeredTraversalPoint,
427        value: &B::Logits,
428        forward: &mut C,
429        context: &B::Context,
430    ) -> Result<B::Logits, E>;
431
432    /// Returns the exact accepted token-id domain for this prediction.
433    fn token_domain(
434        &mut self,
435        prediction: usize,
436        point: LayeredTraversalPoint,
437        forward: &C,
438    ) -> Result<TokenDomain, E>;
439
440    /// Supplies a forced or sampled backend-native token to subsequent units.
441    fn accept(
442        &mut self,
443        prediction: usize,
444        point: LayeredTraversalPoint,
445        token: &B::Token,
446        forward: &mut C,
447        context: &B::Context,
448    ) -> Result<(), E>;
449
450    /// Converts a generic decision failure into the architecture error type.
451    fn decision_error(&mut self, error: SequentialDecisionError<B::Error>) -> E;
452}
453
454/// Adapter that drives sequential decisions from shared layered traversal hooks.
455pub struct SequentialDecisionTraversal<'a, B, S, D, C, E>
456where
457    B: SamplingBackend,
458    S: Sampler<B>,
459    D: SequentialDecisionBoundary<B, C, E>,
460{
461    driver: &'a mut SequentialDecisionDriver<B, S>,
462    boundary: &'a mut D,
463    marker: PhantomData<fn(C) -> E>,
464}
465
466impl<'a, B, S, D, C, E> SequentialDecisionTraversal<'a, B, S, D, C, E>
467where
468    B: SamplingBackend,
469    S: Sampler<B>,
470    D: SequentialDecisionBoundary<B, C, E>,
471{
472    /// Couples a decision driver to one architecture-owned boundary mapping.
473    pub fn new(driver: &'a mut SequentialDecisionDriver<B, S>, boundary: &'a mut D) -> Self {
474        Self {
475            driver,
476            boundary,
477            marker: PhantomData,
478        }
479    }
480
481    fn process(
482        &mut self,
483        point: LayeredTraversalPoint,
484        value: &B::Logits,
485        forward: &mut C,
486        context: &B::Context,
487    ) -> Result<(), E> {
488        let Some(prediction) = self.boundary.prediction_at(point, forward) else {
489            return Ok(());
490        };
491        let logits = self
492            .boundary
493            .logits(prediction, point, value, forward, context)?;
494        let domain = self.boundary.token_domain(prediction, point, forward)?;
495        let token = self
496            .driver
497            .resolve(prediction, &logits, domain, context)
498            .map_err(|error| self.boundary.decision_error(error))?;
499        self.boundary
500            .accept(prediction, point, &token, forward, context)
501    }
502}
503
504impl<NB, B, S, D, C, E> LayeredTraversalHook<NB, C, E>
505    for SequentialDecisionTraversal<'_, B, S, D, C, E>
506where
507    NB: NeuralBackend,
508    B: SamplingBackend<Logits = NB::Tensor, Context = <NB::Tensor as Tensor>::Context>,
509    S: Sampler<B>,
510    D: SequentialDecisionBoundary<B, C, E>,
511{
512    fn before_unit(
513        &mut self,
514        group: usize,
515        index: usize,
516        remaining_units: usize,
517        _value: &mut NB::Tensor,
518        forward: &mut C,
519        context: &<NB::Tensor as Tensor>::Context,
520    ) -> Result<LayeredUnitAction, E> {
521        let point = LayeredTraversalPoint::Unit { group, index };
522        let Some(prediction) = self.boundary.prediction_at(point, forward) else {
523            return Ok(LayeredUnitAction::Execute);
524        };
525        let tail = self
526            .driver
527            .fully_forced_tail_decision(prediction, remaining_units)
528            .map_err(|error| self.boundary.decision_error(error))?;
529        let FullyForcedTailDecision::Skip { predictions } = tail else {
530            return Ok(LayeredUnitAction::Execute);
531        };
532        let mut domains = Vec::with_capacity(predictions);
533        for offset in 0..predictions {
534            let skipped_point = LayeredTraversalPoint::Unit {
535                group,
536                index: index + offset,
537            };
538            let actual = self.boundary.prediction_at(skipped_point, forward);
539            if actual != Some(prediction + offset) {
540                let error = SequentialDecisionError::TailBoundaryMismatch {
541                    expected: prediction + offset,
542                    actual,
543                };
544                return Err(self.boundary.decision_error(error));
545            }
546            domains.push(self.boundary.token_domain(
547                prediction + offset,
548                skipped_point,
549                forward,
550            )?);
551        }
552        let tokens = self
553            .driver
554            .forced_tail_tokens(prediction, predictions, domains, context)
555            .map_err(|error| self.boundary.decision_error(error))?;
556        for (offset, token) in tokens.iter().enumerate() {
557            let skipped_point = LayeredTraversalPoint::Unit {
558                group,
559                index: index + offset,
560            };
561            self.boundary
562                .accept(prediction + offset, skipped_point, token, forward, context)?;
563        }
564        self.driver
565            .commit_forced_tail(prediction, tokens)
566            .map_err(|error| self.boundary.decision_error(error))?;
567        Ok(LayeredUnitAction::SkipRemainingGroup)
568    }
569
570    fn after_unit(
571        &mut self,
572        group: usize,
573        index: usize,
574        value: &mut NB::Tensor,
575        forward: &mut C,
576        context: &<NB::Tensor as Tensor>::Context,
577    ) -> Result<(), E> {
578        self.process(
579            LayeredTraversalPoint::Unit { group, index },
580            value,
581            forward,
582            context,
583        )
584    }
585
586    fn after_group(
587        &mut self,
588        group: usize,
589        value: &mut NB::Tensor,
590        forward: &mut C,
591        context: &<NB::Tensor as Tensor>::Context,
592    ) -> Result<(), E> {
593        self.process(
594            LayeredTraversalPoint::Group { group },
595            value,
596            forward,
597            context,
598        )
599    }
600}
601
602/// Invalid construction of a sequential decision plan or driver.
603#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
604pub enum SequentialDecisionPlanError {
605    /// No target or predictor decisions were declared.
606    #[error("sequential decision plan must contain at least one prediction")]
607    EmptyPlan,
608    /// Sampler state cardinality did not match prediction cardinality.
609    #[error("sequential decision plan has {predictions} predictions but {samplers} samplers")]
610    SamplerCountMismatch {
611        /// Planned prediction count.
612        predictions: usize,
613        /// Supplied sampler count.
614        samplers: usize,
615    },
616    /// Temperature cardinality did not match prediction cardinality.
617    #[error(
618        "sequential decision plan has {predictions} predictions but {temperatures} temperatures"
619    )]
620    TemperatureCountMismatch {
621        /// Planned prediction count.
622        predictions: usize,
623        /// Supplied temperature count.
624        temperatures: usize,
625    },
626    /// One sampling temperature was negative or non-finite.
627    #[error(
628        "sequential decision temperature at prediction {prediction} is invalid (bits {bits:#010x})"
629    )]
630    InvalidTemperature {
631        /// Invalid prediction ordinal.
632        prediction: usize,
633        /// Exact invalid floating-point bits.
634        bits: u32,
635    },
636}
637
638/// Failure while resolving an ordered sequential decision chain.
639#[derive(Debug, thiserror::Error)]
640pub enum SequentialDecisionError<E> {
641    /// The sampler or sampling backend failed.
642    #[error("sequential decision backend failed: {0}")]
643    Backend(E),
644    /// A traversal boundary did not match the next prediction.
645    #[error(
646        "sequential decision expected prediction {expected} of {predictions}, received {actual}"
647    )]
648    OutOfOrder {
649        /// Next required prediction.
650        expected: usize,
651        /// Traversal-supplied prediction.
652        actual: usize,
653        /// Total prediction count.
654        predictions: usize,
655    },
656    /// A requested tail skip was not proven safe by the plan.
657    #[error(
658        "prediction tail beginning at {prediction} with length {count} is not safely skippable"
659    )]
660    InvalidTailSkip {
661        /// First proposed skipped prediction.
662        prediction: usize,
663        /// Proposed skipped prediction count.
664        count: usize,
665    },
666    /// Architecture token-domain cardinality drifted across a skipped tail.
667    #[error(
668        "prediction tail beginning at {prediction} has {actual} token domains, expected {expected}"
669    )]
670    TokenDomainCountMismatch {
671        /// First proposed skipped prediction.
672        prediction: usize,
673        /// Expected domain count.
674        expected: usize,
675        /// Supplied domain count.
676        actual: usize,
677    },
678    /// Architecture boundary mapping disagreed within a proposed skipped tail.
679    #[error("forced-tail boundary expected prediction {expected}, got {actual:?}")]
680    TailBoundaryMismatch {
681        /// Expected prediction ordinal.
682        expected: usize,
683        /// Architecture-reported ordinal.
684        actual: Option<usize>,
685    },
686    /// Traversal completed before resolving every prediction.
687    #[error("sequential decision traversal resolved {resolved} of {predictions} predictions")]
688    Incomplete {
689        /// Resolved prediction count.
690        resolved: usize,
691        /// Planned prediction count.
692        predictions: usize,
693    },
694}
695
696#[cfg(test)]
697mod tests {
698    use super::*;
699    use crate::PenaltyConfig;
700    use eredu_core::TokenFilter;
701
702    struct Backend;
703
704    impl SamplingBackend for Backend {
705        type Logits = i32;
706        type Token = i32;
707        type RandomState = i32;
708        type Context = ();
709        type Error = String;
710
711        fn error(message: String) -> Self::Error {
712            message
713        }
714
715        fn validate_token(
716            token: &Self::Token,
717            domain: TokenDomain,
718            _: &Self::Context,
719        ) -> Result<Self::Token, Self::Error> {
720            usize::try_from(*token)
721                .ok()
722                .filter(|token| *token < domain.cardinality())
723                .map(|_| *token)
724                .ok_or_else(|| "token is outside its decision domain".into())
725        }
726
727        fn scale_temperature(
728            logits: &Self::Logits,
729            _: f32,
730            _: &Self::Context,
731        ) -> Result<Self::Logits, Self::Error> {
732            Ok(*logits)
733        }
734
735        fn apply_penalties(
736            logits: &Self::Logits,
737            _: &[u32],
738            _: PenaltyConfig,
739            _: &Self::Context,
740        ) -> Result<Self::Logits, Self::Error> {
741            Ok(*logits)
742        }
743
744        fn apply_top_k(
745            logits: Self::Logits,
746            _: i32,
747            _: &Self::Context,
748        ) -> Result<Self::Logits, Self::Error> {
749            Ok(logits)
750        }
751
752        fn apply_top_p(
753            logits: Self::Logits,
754            _: f32,
755            _: &Self::Context,
756        ) -> Result<Self::Logits, Self::Error> {
757            Ok(logits)
758        }
759
760        fn apply_min_p(
761            logits: Self::Logits,
762            _: f32,
763            _: &Self::Context,
764        ) -> Result<Self::Logits, Self::Error> {
765            Ok(logits)
766        }
767
768        fn apply_token_filter(
769            logits: &Self::Logits,
770            _: &TokenFilter,
771            _: &Self::Context,
772        ) -> Result<Self::Logits, Self::Error> {
773            Ok(*logits)
774        }
775
776        fn apply_mirostat(
777            logits: &Self::Logits,
778            _: &[u32],
779            _: PenaltyConfig,
780            _: f32,
781            _: f32,
782            _: &Self::Context,
783        ) -> Result<Self::Logits, Self::Error> {
784            Ok(*logits)
785        }
786
787        fn sample_raw(
788            logits: &Self::Logits,
789            _: f32,
790            random: Option<&mut Self::RandomState>,
791            _: &Self::Context,
792        ) -> Result<Self::Token, Self::Error> {
793            if let Some(random) = random {
794                *random += 1;
795            }
796            Ok(*logits)
797        }
798
799        fn sample_processed(
800            logits: &Self::Logits,
801            temperature: f32,
802            random: Option<&mut Self::RandomState>,
803            context: &Self::Context,
804        ) -> Result<Self::Token, Self::Error> {
805            Self::sample_raw(logits, temperature, random, context)
806        }
807
808        fn token_id(token: &Self::Token, _: &Self::Context) -> Result<u32, Self::Error> {
809            u32::try_from(*token).map_err(|error| error.to_string())
810        }
811
812        fn token_probability(
813            _: &Self::Logits,
814            _: u32,
815            _: &Self::Context,
816        ) -> Result<f32, Self::Error> {
817            Ok(1.0)
818        }
819    }
820
821    #[derive(Clone)]
822    struct OffsetSampler(i32);
823
824    impl Sampler<Backend> for OffsetSampler {
825        fn sample(
826            &mut self,
827            logits: &i32,
828            _: f32,
829            random: Option<&mut i32>,
830            _: &(),
831        ) -> Result<i32, String> {
832            if let Some(random) = random {
833                *random += 1;
834            }
835            Ok(*logits + self.0)
836        }
837    }
838
839    #[test]
840    fn teacher_forcing_sampling_masks_and_diagnostics_share_one_driver() {
841        let plan = SequentialDecisionPlan::new(
842            [
843                PredictionDirective::Force(7),
844                PredictionDirective::Sample,
845                PredictionDirective::Force(9),
846            ],
847            true,
848            true,
849        )
850        .unwrap();
851        assert_eq!(plan.mode(), SequentialDecisionMode::PartiallyForced);
852        assert_eq!(plan.forcing_mask().collect::<Vec<_>>(), [true, false, true]);
853        let mut driver = SequentialDecisionDriver::<Backend, _>::new(
854            plan,
855            vec![OffsetSampler(100), OffsetSampler(10), OffsetSampler(100)],
856            vec![0.0; 3],
857            Some(4),
858        )
859        .unwrap();
860
861        let domain = TokenDomain::new(100);
862        assert_eq!(driver.resolve(0, &1, domain, &()).unwrap(), 7);
863        assert_eq!(driver.resolve(1, &2, domain, &()).unwrap(), 12);
864        assert_eq!(driver.resolve(2, &3, domain, &()).unwrap(), 9);
865        driver.finish().unwrap();
866        assert_eq!(driver.random_state(), Some(&5));
867        assert_eq!(
868            driver
869                .decisions()
870                .iter()
871                .map(|decision| (decision.source(), *decision.token()))
872                .collect::<Vec<_>>(),
873            [
874                (SequentialDecisionSource::Forced, 7),
875                (SequentialDecisionSource::Sampled, 12),
876                (SequentialDecisionSource::Forced, 9),
877            ]
878        );
879        assert_eq!(
880            driver
881                .diagnostics()
882                .iter()
883                .map(|diagnostic| (diagnostic.prediction(), *diagnostic.logits()))
884                .collect::<Vec<_>>(),
885            [(0, 1), (1, 2), (2, 3)]
886        );
887    }
888
889    #[test]
890    fn fully_forced_tail_skip_requires_exact_cardinality_and_no_diagnostics() {
891        let plan = SequentialDecisionPlan::new(
892            [
893                PredictionDirective::Sample,
894                PredictionDirective::Force(8),
895                PredictionDirective::Force(9),
896            ],
897            false,
898            true,
899        )
900        .unwrap();
901        let mut driver = SequentialDecisionDriver::<Backend, _>::new(
902            plan,
903            vec![OffsetSampler(1); 3],
904            vec![0.0; 3],
905            None,
906        )
907        .unwrap();
908        driver.resolve(0, &4, TokenDomain::new(100), &()).unwrap();
909        assert_eq!(
910            driver.fully_forced_tail_decision(1, 1).unwrap(),
911            FullyForcedTailDecision::Execute
912        );
913        assert_eq!(
914            driver.fully_forced_tail_decision(1, 2).unwrap(),
915            FullyForcedTailDecision::Skip { predictions: 2 }
916        );
917        let tokens = driver
918            .forced_tail_tokens(1, 2, [TokenDomain::new(100); 2], &())
919            .unwrap();
920        driver.commit_forced_tail(1, tokens).unwrap();
921        driver.finish().unwrap();
922        assert_eq!(
923            driver
924                .decisions()
925                .iter()
926                .map(SequentialDecision::source)
927                .collect::<Vec<_>>(),
928            [
929                SequentialDecisionSource::Sampled,
930                SequentialDecisionSource::ForcedTailSkipped,
931                SequentialDecisionSource::ForcedTailSkipped,
932            ]
933        );
934        assert!(driver.diagnostics().is_empty());
935
936        let diagnostic_plan = SequentialDecisionPlan::new(
937            [PredictionDirective::Force(1), PredictionDirective::Force(2)],
938            true,
939            true,
940        )
941        .unwrap();
942        let diagnostic = SequentialDecisionDriver::<Backend, _>::new(
943            diagnostic_plan,
944            vec![OffsetSampler(0); 2],
945            vec![0.0; 2],
946            None,
947        )
948        .unwrap();
949        assert_eq!(
950            diagnostic.fully_forced_tail_decision(0, 2).unwrap(),
951            FullyForcedTailDecision::Execute
952        );
953    }
954
955    #[test]
956    fn driver_rejects_cardinality_temperature_and_order_drift() {
957        let plan =
958            SequentialDecisionPlan::new([PredictionDirective::Sample], false, false).unwrap();
959        assert!(matches!(
960            SequentialDecisionDriver::<Backend, _>::new(
961                plan.clone(),
962                Vec::<OffsetSampler>::new(),
963                vec![0.0],
964                None
965            ),
966            Err(SequentialDecisionPlanError::SamplerCountMismatch { .. })
967        ));
968        assert!(matches!(
969            SequentialDecisionDriver::<Backend, _>::new(
970                plan.clone(),
971                vec![OffsetSampler(0)],
972                vec![f32::NAN],
973                None
974            ),
975            Err(SequentialDecisionPlanError::InvalidTemperature { .. })
976        ));
977        let mut driver = SequentialDecisionDriver::<Backend, _>::new(
978            plan,
979            vec![OffsetSampler(0)],
980            vec![0.0],
981            None,
982        )
983        .unwrap();
984        assert!(matches!(
985            driver.resolve(1, &0, TokenDomain::new(100), &()),
986            Err(SequentialDecisionError::OutOfOrder { .. })
987        ));
988    }
989
990    #[test]
991    fn token_domains_reject_sampled_executed_forcing_and_skipped_forcing_before_commit() {
992        let domain = TokenDomain::new(10);
993
994        let forced_plan =
995            SequentialDecisionPlan::new([PredictionDirective::Force(10)], false, false).unwrap();
996        let mut forced = SequentialDecisionDriver::<Backend, _>::new(
997            forced_plan,
998            vec![OffsetSampler(0)],
999            vec![0.0],
1000            None,
1001        )
1002        .unwrap();
1003        assert!(matches!(
1004            forced.resolve(0, &0, domain, &()),
1005            Err(SequentialDecisionError::Backend(_))
1006        ));
1007        assert!(forced.decisions().is_empty());
1008
1009        let sampled_plan =
1010            SequentialDecisionPlan::new([PredictionDirective::Sample], false, false).unwrap();
1011        let mut sampled = SequentialDecisionDriver::<Backend, _>::new(
1012            sampled_plan,
1013            vec![OffsetSampler(10)],
1014            vec![0.0],
1015            Some(4),
1016        )
1017        .unwrap();
1018        assert!(matches!(
1019            sampled.resolve(0, &0, domain, &()),
1020            Err(SequentialDecisionError::Backend(_))
1021        ));
1022        assert!(sampled.decisions().is_empty());
1023
1024        let tail_plan = SequentialDecisionPlan::new(
1025            [
1026                PredictionDirective::Force(1),
1027                PredictionDirective::Force(10),
1028            ],
1029            false,
1030            true,
1031        )
1032        .unwrap();
1033        let tail = SequentialDecisionDriver::<Backend, _>::new(
1034            tail_plan,
1035            vec![OffsetSampler(0); 2],
1036            vec![0.0; 2],
1037            None,
1038        )
1039        .unwrap();
1040        assert!(matches!(
1041            tail.forced_tail_tokens(0, 2, [domain; 2], &()),
1042            Err(SequentialDecisionError::Backend(_))
1043        ));
1044        assert!(tail.decisions().is_empty());
1045    }
1046}
eredu_runtime/decision.rs

eredu_runtime/
decision.rs
