tract_core/ops/nn/
clamped_swiglu.rs

1use crate::internal::*;
2
3/// Two-input GPT-OSS clamped SwiGLU, without broadcasting.
4///
5/// For `g = min(gate, limit)` and `u = clamp(up, -limit, limit)`, computes
6/// `(u + 1) * g * sigmoid(alpha * g)`. Inputs must have equal shapes and
7/// f16 or f32 elements; computation and output are f32. Both parameters must
8/// be finite, and `limit` must be positive.
9#[derive(Clone, Debug, PartialEq)]
10pub struct ClampedSwiGlu {
11    pub alpha: f32,
12    pub limit: f32,
13}
14
15impl Eq for ClampedSwiGlu {}
16
17impl ClampedSwiGlu {
18    fn validate(&self) -> TractResult<()> {
19        ensure!(self.alpha.is_finite(), "ClampedSwiGlu alpha must be finite");
20        ensure!(
21            self.limit.is_finite() && self.limit > 0.0,
22            "ClampedSwiGlu limit must be finite and positive"
23        );
24        Ok(())
25    }
26}
27
28impl Op for ClampedSwiGlu {
29    fn name(&self) -> StaticName {
30        "ClampedSwiGlu".into()
31    }
32    op_as_typed_op!();
33}
34
35impl EvalOp for ClampedSwiGlu {
36    op_out_of_plan!();
37
38    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
39        self.validate()?;
40        ensure!(inputs.len() == 2, "ClampedSwiGlu expects gate and up inputs");
41        ensure!(inputs[0].shape() == inputs[1].shape(), "ClampedSwiGlu gate/up shape mismatch");
42        ensure!(
43            inputs.iter().all(|t| matches!(t.datum_type(), DatumType::F16 | DatumType::F32)),
44            "ClampedSwiGlu inputs must be f16 or f32"
45        );
46        let mut output = Tensor::zero_dt(f32::datum_type(), inputs[0].shape())?;
47        if output.len() == 0 {
48            return Ok(tvec![output.into_tvalue()]);
49        }
50        let (gate, up) = args_2!(inputs);
51        let gate = gate.cast_to::<f32>()?;
52        let up = up.cast_to::<f32>()?;
53        let gate_ram = gate.try_as_plain_ram()?;
54        let up_ram = up.try_as_plain_ram()?;
55        let mut output_ram = output.try_as_plain_ram_mut()?;
56        let gate = gate_ram.as_slice::<f32>()?;
57        let up = up_ram.as_slice::<f32>()?;
58        let output_slice = output_ram.as_slice_mut::<f32>()?;
59        for ((out, &gate), &up) in output_slice.iter_mut().zip(gate).zip(up) {
60            let gate = gate.min(self.limit);
61            let up = up.clamp(-self.limit, self.limit);
62            let glu = gate / (1.0 + (-self.alpha * gate).exp());
63            *out = (up + 1.0) * glu;
64        }
65        Ok(tvec![output.into_tvalue()])
66    }
67}
68
69impl TypedOp for ClampedSwiGlu {
70    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
71        self.validate()?;
72        ensure!(inputs.len() == 2, "ClampedSwiGlu expects gate and up inputs");
73        ensure!(inputs[0].shape == inputs[1].shape, "ClampedSwiGlu gate/up shape mismatch");
74        ensure!(
75            inputs.iter().all(|f| matches!(f.datum_type, DatumType::F16 | DatumType::F32)),
76            "ClampedSwiGlu inputs must be f16 or f32"
77        );
78        Ok(tvec![f32::datum_type().fact(inputs[0].shape.clone())])
79    }
80    as_op!();
81}
tract_core/ops/nn/clamped_swiglu.rs

tract_core/ops/nn/
clamped_swiglu.rs
