eredu_runtime/intervention/
activation.rs

1use super::*;
2
3/// Executes the shared activation recipe. Direct callers receive the same runtime
4/// validation as admitted runs; no primitive receives the public action enum.
5pub fn apply_activation<B: InterventionBackend>(
6    backend: &mut B,
7    input: &B::Tensor,
8    action: &InterventionAction,
9    slice: &ResolvedCaptureSlice,
10) -> Result<B::Tensor, CaptureExecutionError<B::Error>> {
11    use CaptureExecutionError::Backend;
12    let source = backend.shape(input).map_err(Backend)?;
13    let dtype = backend.intervention_dtype(input).map_err(Backend)?;
14    let rank = source.len();
15    if [&slice.starts, &slice.ends, &slice.strides, &slice.shape]
16        .iter()
17        .any(|v| v.len() != rank)
18    {
19        return Err(CaptureError::Invalid("activation slice rank mismatch".into()).into());
20    }
21    for (axis, extent) in source.iter().enumerate() {
22        let (start, end, stride) = (slice.starts[axis], slice.ends[axis], slice.strides[axis]);
23        if stride == 0
24            || start >= end
25            || end > *extent
26            || (end - start).div_ceil(stride) != slice.shape[axis]
27        {
28            return Err(CaptureError::Invalid("invalid activation region".into()).into());
29        }
30    }
31    action.validate_activation_region(dtype, &slice.shape)?;
32    if matches!(action, InterventionAction::MaskLogits { .. })
33        && (slice.starts[rank - 1] != 0
34            || slice.ends[rank - 1] != source[rank - 1]
35            || slice.strides[rank - 1] != 1)
36    {
37        return Err(CaptureError::Invalid(
38            "logit masks require the complete vocabulary axis".into(),
39        )
40        .into());
41    }
42    backend.validate_intervention_geometry(&source, slice)?;
43    let selected = backend.select_region(input, slice).map_err(Backend)?;
44    validate_value(backend, &selected, &slice.shape, dtype)?;
45    let replacement = match action {
46        InterventionAction::Zero { .. } => backend.zeros(&slice.shape, dtype).map_err(Backend)?,
47        InterventionAction::Scale { factor, .. } => {
48            backend.scale(&selected, *factor).map_err(Backend)?
49        }
50        InterventionAction::Mask { keep, .. } => {
51            backend.fill_masked(&selected, keep, 0.0).map_err(Backend)?
52        }
53        InterventionAction::Replace { tensor } => {
54            backend.realize_tensor(tensor).map_err(Backend)?
55        }
56        InterventionAction::Add { tensor } => {
57            let delta = backend.realize_tensor(tensor).map_err(Backend)?;
58            validate_value(backend, &delta, &slice.shape, dtype)?;
59            backend.add(&selected, &delta).map_err(Backend)?
60        }
61        InterventionAction::MaskLogits { token_ids, .. } => backend
62            .fill_columns(&selected, token_ids, f32::NEG_INFINITY)
63            .map_err(Backend)?,
64        _ => {
65            return Err(
66                CaptureError::Invalid("routing action cannot replace an activation".into()).into(),
67            )
68        }
69    };
70    validate_value(backend, &replacement, &slice.shape, dtype)?;
71    let output = backend
72        .update_region(input, slice, &replacement)
73        .map_err(Backend)?;
74    validate_value(backend, &output, &source, dtype)?;
75    Ok(output)
76}
77
78fn validate_value<B: InterventionBackend>(
79    backend: &B,
80    value: &B::Tensor,
81    shape: &[u64],
82    dtype: InterventionDtype,
83) -> Result<(), CaptureExecutionError<B::Error>> {
84    if backend
85        .shape(value)
86        .map_err(CaptureExecutionError::Backend)?
87        != shape
88        || backend
89            .intervention_dtype(value)
90            .map_err(CaptureExecutionError::Backend)?
91            != dtype
92    {
93        return Err(CaptureError::Invalid(
94            "native activation primitive changed shape or dtype".into(),
95        )
96        .into());
97    }
98    Ok(())
99}
eredu_runtime/intervention/activation.rs

eredu_runtime/intervention/
activation.rs
