candle_graph/
op_semantics.rs

1//! Transfer rules for candle Tensor ops and known candle-nn helpers.
2//!
3//! Rules are deliberate and incomplete by design: anything not listed returns
4//! [`OpEffect::Unknown`] rather than inventing behaviour.
5
6use serde::Serialize;
7
8/// The Candle release whose version-sensitive implementation details back this catalog.
9pub const AUDITED_CANDLE_VERSION: &str = "0.11.0";
10
11pub fn is_audited_candle_version(version: &str) -> bool {
12    version == AUDITED_CANDLE_VERSION
13}
14
15/// Return the shared audited version only when candle-core and candle-nn resolve consistently.
16///
17/// A mismatched pair is not a sound basis for candle-nn rules because those implementations call
18/// directly into candle-core's operation and autograd APIs.
19pub fn matched_candle_version<'a>(
20    candle_core_version: Option<&'a str>,
21    candle_nn_version: Option<&'a str>,
22) -> Option<&'a str> {
23    match (candle_core_version, candle_nn_version) {
24        (Some(core), Some(nn)) if core == nn && is_audited_candle_version(nn) => Some(nn),
25        _ => None,
26    }
27}
28
29/// Abstract dtype carried on expression nodes. Coarse by intent.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize)]
31#[serde(rename_all = "snake_case")]
32pub enum AbstractDtype {
33    F64,
34    F32,
35    F16,
36    Bf16,
37    I16,
38    I32,
39    I64,
40    U32,
41    U8,
42    F8E4M3,
43    F6E2M3,
44    F6E3M2,
45    F4,
46    F8E8M0,
47    /// Could not be determined statically.
48    Unknown,
49}
50
51impl AbstractDtype {
52    pub fn parse(text: &str) -> Self {
53        let t = text.trim();
54        let t = t.rsplit("::").next().unwrap_or(t);
55        match t {
56            "F64" | "f64" => Self::F64,
57            "F32" | "f32" => Self::F32,
58            "F16" | "f16" => Self::F16,
59            "BF16" | "Bf16" | "bf16" => Self::Bf16,
60            "I16" | "i16" => Self::I16,
61            "I32" | "i32" => Self::I32,
62            "I64" | "i64" => Self::I64,
63            "U32" | "u32" => Self::U32,
64            "U8" | "u8" => Self::U8,
65            "F8E4M3" => Self::F8E4M3,
66            "F6E2M3" => Self::F6E2M3,
67            "F6E3M2" => Self::F6E3M2,
68            "F4" => Self::F4,
69            "F8E8M0" => Self::F8E8M0,
70            _ => Self::Unknown,
71        }
72    }
73
74    pub fn is_known(self) -> bool {
75        !matches!(self, Self::Unknown)
76    }
77}
78
79impl std::fmt::Display for AbstractDtype {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        match self {
82            Self::F64 => write!(f, "F64"),
83            Self::F32 => write!(f, "F32"),
84            Self::F16 => write!(f, "F16"),
85            Self::Bf16 => write!(f, "BF16"),
86            Self::I16 => write!(f, "I16"),
87            Self::I32 => write!(f, "I32"),
88            Self::I64 => write!(f, "I64"),
89            Self::U32 => write!(f, "U32"),
90            Self::U8 => write!(f, "U8"),
91            Self::F8E4M3 => write!(f, "F8E4M3"),
92            Self::F6E2M3 => write!(f, "F6E2M3"),
93            Self::F6E3M2 => write!(f, "F6E3M2"),
94            Self::F4 => write!(f, "F4"),
95            Self::F8E8M0 => write!(f, "F8E8M0"),
96            Self::Unknown => write!(f, "Unknown"),
97        }
98    }
99}
100
101/// How an op affects gradient connectivity relative to its tensor inputs.
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
103#[serde(rename_all = "snake_case")]
104pub enum GradFlow {
105    /// Gradients flow through all tensor operands (ordinary differentiable op).
106    Propagates,
107    /// Explicitly cuts the autograd graph (`detach`, `apply_op*_no_bwd`, …).
108    Severs,
109    /// Gradient exists only under layout/contiguity assumptions the analyzer cannot prove.
110    LayoutDependent,
111    /// Not enough information.
112    Unknown,
113}
114
115/// How the result dtype relates to operand dtypes.
116#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
117#[serde(rename_all = "snake_case")]
118pub enum DtypeRule {
119    /// Result dtype equals the (sole / first) tensor operand.
120    Preserve,
121    /// All tensor operands must share a dtype; result is that dtype. Mismatch is a conflict.
122    SameAsInputs,
123    /// Result dtype is taken from an explicit argument (e.g. `to_dtype(DType::F32)`).
124    Explicit,
125    /// Result has a fixed Candle dtype independent of its operands.
126    Fixed(AbstractDtype),
127    /// Result dtype cannot be stated.
128    Unknown,
129}
130
131/// Float range of an op's result after rounding — not the mathematical range.
132#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
133#[serde(rename_all = "snake_case")]
134pub enum NumericDomain {
135    Real,
136    NonNegative,
137    /// Mathematically open; attains `0.0` and `1.0` exactly after f32 rounding.
138    SaturatingUnit,
139    /// Provably strictly inside its bounds (e.g. after an epsilon-guard `affine`).
140    StrictlyPositive,
141    Unknown,
142}
143
144/// Precondition an op places on its primary tensor operand.
145#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
146#[serde(rename_all = "snake_case")]
147pub enum DomainRequirement {
148    StrictlyPositive,
149    NonZero,
150    NonNegative,
151    None,
152}
153
154/// Measured f32 saturation thresholds for candle-nn 0.11.0 `sigmoid` =
155/// `(exp(-v) + 1).recip()` (`candle-nn-0.11.0/src/ops.rs:59-60`).
156pub const SIGMOID_F32_UPPER_SATURATION: f32 = 16.6355;
157pub const SIGMOID_F32_LOWER_SATURATION: f32 = -88.7228;
158
159/// One step in an audited library-op body. Step indices refer to earlier atoms in the same body.
160#[derive(Debug, Clone, Copy, PartialEq)]
161pub enum BodyAtom {
162    /// Call-site argument (`0` = first tensor operand).
163    Arg(usize),
164    /// Identity view that records an API-level domain fact (e.g. BCE targets ∈ [0, 1]).
165    Assume {
166        src: u16,
167        domain: NumericDomain,
168    },
169    Unary {
170        op: &'static str,
171        src: u16,
172    },
173    Binary {
174        op: &'static str,
175        left: u16,
176        right: u16,
177    },
178    /// `tensor.affine(mul, add)` with proven literal coefficients.
179    Affine {
180        src: u16,
181        mul: f64,
182        add: f64,
183    },
184}
185
186/// Expandable body of a candle-nn (or similar) helper, audited for one Candle release.
187#[derive(Debug, Clone, Copy, PartialEq)]
188pub struct LibraryBody {
189    pub cite: &'static str,
190    pub steps: &'static [BodyAtom],
191}
192
193/// candle-nn 0.11.0 `binary_cross_entropy_with_logit` (`loss.rs:64-74`):
194/// `sigmoid` → `log` / `(1-p).log` → multiply by target coefficients.
195const BCE_WITH_LOGITS_011: LibraryBody = LibraryBody {
196    cite: "candle-nn-0.11.0/src/loss.rs:64-74",
197    steps: &[
198        BodyAtom::Arg(0), // 0: logits
199        BodyAtom::Arg(1), // 1: target
200        // candle-nn documents targets as unit labels / probabilities.
201        BodyAtom::Assume {
202            src: 1,
203            domain: NumericDomain::SaturatingUnit,
204        }, // 2: t
205        BodyAtom::Unary {
206            op: "sigmoid",
207            src: 0,
208        }, // 3: p
209        BodyAtom::Unary { op: "log", src: 3 }, // 4: log(p)
210        BodyAtom::Binary {
211            op: "mul",
212            left: 2,
213            right: 4,
214        }, // 5: t * log(p)
215        BodyAtom::Affine {
216            src: 3,
217            mul: -1.0,
218            add: 1.0,
219        }, // 6: 1 - p
220        BodyAtom::Unary { op: "log", src: 6 }, // 7: log(1-p)
221        BodyAtom::Affine {
222            src: 2,
223            mul: -1.0,
224            add: 1.0,
225        }, // 8: 1 - t
226        BodyAtom::Binary {
227            op: "mul",
228            left: 8,
229            right: 7,
230        }, // 9: (1-t) * log(1-p) — 0 * -inf when saturated
231        BodyAtom::Binary {
232            op: "add",
233            left: 5,
234            right: 9,
235        }, // 10
236    ],
237};
238
239/// Look up an expandable body for `op` under an audited Candle version.
240///
241/// Bodies are transfer recipes, not precomputed “unsafe” verdicts. The dataflow domain pass
242/// decides whether the expanded composition is hazardous.
243pub fn library_body(op: &str, candle_nn_version: Option<&str>) -> Option<&'static LibraryBody> {
244    if !candle_nn_version.is_some_and(is_audited_candle_version) {
245        return None;
246    }
247    let name = op.rsplit("::").next().unwrap_or(op);
248    match name {
249        "binary_cross_entropy_with_logit" => Some(&BCE_WITH_LOGITS_011),
250        _ => None,
251    }
252}
253
254#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
255pub struct OpEffect {
256    pub name: String,
257    pub dtype: DtypeRule,
258    pub grad: GradFlow,
259    /// True when this call is treated as a scalar loss sink for connectivity queries.
260    pub is_loss: bool,
261    /// Optional note carried into diagnostics (citation / reason).
262    pub note: Option<&'static str>,
263    /// Result domain after float rounding. `Unknown` when not audited.
264    pub domain: NumericDomain,
265    /// Requirement on the primary tensor operand.
266    pub requires: DomainRequirement,
267}
268
269impl OpEffect {
270    fn known(
271        name: &str,
272        dtype: DtypeRule,
273        grad: GradFlow,
274        is_loss: bool,
275        note: Option<&'static str>,
276    ) -> Self {
277        Self {
278            name: name.to_string(),
279            dtype,
280            grad,
281            is_loss,
282            note,
283            domain: NumericDomain::Real,
284            requires: DomainRequirement::None,
285        }
286    }
287
288    fn with_domain(mut self, domain: NumericDomain) -> Self {
289        self.domain = domain;
290        self
291    }
292
293    fn with_requires(mut self, requires: DomainRequirement) -> Self {
294        self.requires = requires;
295        self
296    }
297
298    pub fn unknown(name: &str) -> Self {
299        Self {
300            name: name.to_string(),
301            dtype: DtypeRule::Unknown,
302            grad: GradFlow::Unknown,
303            is_loss: false,
304            note: Some("no transfer rule; left Unknown"),
305            domain: NumericDomain::Unknown,
306            requires: DomainRequirement::None,
307        }
308    }
309
310    pub fn domain_rule_label(&self) -> String {
311        match self.domain {
312            NumericDomain::Real => "real".into(),
313            NumericDomain::NonNegative => "non_negative".into(),
314            NumericDomain::SaturatingUnit => format!(
315                "saturating_unit(f32_upper={SIGMOID_F32_UPPER_SATURATION},f32_lower={SIGMOID_F32_LOWER_SATURATION})"
316            ),
317            NumericDomain::StrictlyPositive => "strictly_positive".into(),
318            NumericDomain::Unknown => "unknown".into(),
319        }
320    }
321}
322
323/// Strength of a domain-requirement failure.
324#[derive(Debug, Clone, Copy, PartialEq, Eq)]
325pub enum DomainViolationConfidence {
326    Proven,
327    Unknown,
328}
329
330/// Whether `domain` fails to discharge `requires`.
331///
332/// Returns `Some(Proven)` for a catalog-known violation, `Some(Unknown)` when the producer
333/// domain is unknown, and `None` when the requirement is met.
334pub fn domain_violation(
335    requires: DomainRequirement,
336    domain: NumericDomain,
337) -> Option<DomainViolationConfidence> {
338    match requires {
339        DomainRequirement::None => None,
340        DomainRequirement::StrictlyPositive => match domain {
341            NumericDomain::StrictlyPositive => None,
342            NumericDomain::Unknown => Some(DomainViolationConfidence::Unknown),
343            NumericDomain::Real | NumericDomain::NonNegative | NumericDomain::SaturatingUnit => {
344                Some(DomainViolationConfidence::Proven)
345            }
346        },
347        DomainRequirement::NonZero => match domain {
348            NumericDomain::StrictlyPositive => None,
349            NumericDomain::Unknown => Some(DomainViolationConfidence::Unknown),
350            NumericDomain::Real | NumericDomain::NonNegative | NumericDomain::SaturatingUnit => {
351                Some(DomainViolationConfidence::Proven)
352            }
353        },
354        DomainRequirement::NonNegative => match domain {
355            NumericDomain::NonNegative
356            | NumericDomain::StrictlyPositive
357            | NumericDomain::SaturatingUnit => None,
358            NumericDomain::Unknown => Some(DomainViolationConfidence::Unknown),
359            NumericDomain::Real => Some(DomainViolationConfidence::Proven),
360        },
361    }
362}
363
364/// Join domains through a dtype-preserving unary/binary where both sides contribute.
365pub fn join_domain(left: NumericDomain, right: NumericDomain) -> NumericDomain {
366    use NumericDomain::*;
367    match (left, right) {
368        (Unknown, _) | (_, Unknown) => Unknown,
369        (StrictlyPositive, StrictlyPositive) => StrictlyPositive,
370        (StrictlyPositive, NonNegative) | (NonNegative, StrictlyPositive) => NonNegative,
371        (NonNegative, NonNegative) => NonNegative,
372        (SaturatingUnit, SaturatingUnit) => SaturatingUnit,
373        (SaturatingUnit, NonNegative) | (NonNegative, SaturatingUnit) => NonNegative,
374        (SaturatingUnit, StrictlyPositive) | (StrictlyPositive, SaturatingUnit) => NonNegative,
375        _ => Real,
376    }
377}
378
379/// Transfer for `tensor.affine(mul, add)` with proven literal coefficients.
380///
381/// Only the epsilon-guard form `mul > 0 && add > 0` on a non-negative / unit operand
382/// discharges `StrictlyPositive`. Reflections such as `affine(-1, 1)` (`1 - p`) keep a
383/// zero-attaining domain.
384pub fn affine_domain(operand: NumericDomain, mul: Option<f64>, add: Option<f64>) -> NumericDomain {
385    match (mul, add) {
386        (Some(m), Some(a)) if m > 0.0 && a > 0.0 => match operand {
387            NumericDomain::SaturatingUnit
388            | NumericDomain::NonNegative
389            | NumericDomain::StrictlyPositive => NumericDomain::StrictlyPositive,
390            NumericDomain::Unknown => NumericDomain::Unknown,
391            NumericDomain::Real => NumericDomain::Real,
392        },
393        _ => match operand {
394            // 1 - p over a unit interval still attains 0 and 1.
395            NumericDomain::SaturatingUnit => NumericDomain::SaturatingUnit,
396            NumericDomain::StrictlyPositive => NumericDomain::Real,
397            NumericDomain::NonNegative => NumericDomain::Real,
398            NumericDomain::Unknown => NumericDomain::Unknown,
399            NumericDomain::Real => NumericDomain::Real,
400        },
401    }
402}
403
404pub fn domain_includes_zero(domain: NumericDomain) -> bool {
405    matches!(
406        domain,
407        NumericDomain::NonNegative | NumericDomain::SaturatingUnit | NumericDomain::Real
408    )
409}
410
411/// Last path segment of a call, e.g. `broadcast_add` or `cross_entropy`.
412pub fn lookup(op: &str) -> OpEffect {
413    // Without Cargo evidence the version-sensitive answer is unknown. Callers that resolved the
414    // target crate must use `lookup_for`.
415    lookup_for(op, None)
416}
417
418/// Look up an operation against an audited Candle version.
419///
420/// Generic Tensor algebra is stable across the supported catalog. Rules tied to Candle's
421/// custom `*_no_bwd` implementations are only asserted for versions whose source was audited;
422/// other versions stay `Unknown` rather than inheriting a potentially stale gradient claim.
423pub fn lookup_for(op: &str, candle_nn_version: Option<&str>) -> OpEffect {
424    let name = op.rsplit("::").next().unwrap_or(op);
425    let audited_no_bwd = matches!(
426        name,
427        "softmax_last_dim"
428            | "rms_norm"
429            | "layer_norm"
430            | "sdpa"
431            | "rope"
432            | "rope_i"
433            | "rope_thd"
434            | "flash_attn"
435            | "flash_attn_varlen_cpu"
436            | "flash_attn_varlen_unfused"
437            | "run_flash_attn_cpu"
438    );
439    if audited_no_bwd && !candle_nn_version.is_some_and(is_audited_candle_version) {
440        return OpEffect::unknown(name);
441    }
442    match name {
443        // --- explicit dtype ---
444        "to_dtype" => OpEffect::known(
445            name,
446            DtypeRule::Explicit,
447            GradFlow::Propagates,
448            false,
449            Some("candle Tensor::to_dtype"),
450        ),
451
452        // --- same-dtype binaries (conflict when operands disagree) ---
453        "add" | "sub" | "mul" | "maximum" | "minimum" => OpEffect::known(
454            name,
455            DtypeRule::SameAsInputs,
456            GradFlow::Propagates,
457            false,
458            Some("candle binary op; operands must share dtype"),
459        ),
460        "div" => OpEffect::known(
461            name,
462            DtypeRule::SameAsInputs,
463            GradFlow::Propagates,
464            false,
465            Some("candle binary op; operands must share dtype"),
466        )
467        .with_requires(DomainRequirement::NonZero),
468        "broadcast_add" | "broadcast_sub" | "broadcast_mul" | "broadcast_maximum"
469        | "broadcast_minimum" | "broadcast_pow" => OpEffect::known(
470            name,
471            DtypeRule::SameAsInputs,
472            GradFlow::Propagates,
473            false,
474            Some("candle broadcast binary; operands must share dtype"),
475        ),
476        "broadcast_div" => OpEffect::known(
477            name,
478            DtypeRule::SameAsInputs,
479            GradFlow::Propagates,
480            false,
481            Some("candle broadcast binary; operands must share dtype"),
482        )
483        .with_requires(DomainRequirement::NonZero),
484        "matmul" | "broadcast_matmul" => OpEffect::known(
485            name,
486            DtypeRule::SameAsInputs,
487            GradFlow::Propagates,
488            false,
489            Some("candle matmul; operands must share dtype"),
490        ),
491
492        // --- dtype-preserving unary math ---
493        "neg" | "sin" | "cos" | "tanh" | "gelu" | "silu" | "erf" => OpEffect::known(
494            name,
495            DtypeRule::Preserve,
496            GradFlow::Propagates,
497            false,
498            Some("dtype-preserving unary"),
499        ),
500        "abs" | "sqr" | "relu" => OpEffect::known(
501            name,
502            DtypeRule::Preserve,
503            GradFlow::Propagates,
504            false,
505            Some("dtype-preserving unary"),
506        )
507        .with_domain(NumericDomain::NonNegative),
508        "exp" => OpEffect::known(
509            name,
510            DtypeRule::Preserve,
511            GradFlow::Propagates,
512            false,
513            Some("dtype-preserving unary; f32 underflows to exactly 0.0"),
514        )
515        .with_domain(NumericDomain::NonNegative),
516        "sqrt" => OpEffect::known(
517            name,
518            DtypeRule::Preserve,
519            GradFlow::Propagates,
520            false,
521            Some("dtype-preserving unary"),
522        )
523        .with_domain(NumericDomain::NonNegative)
524        .with_requires(DomainRequirement::NonNegative),
525        "log" => OpEffect::known(
526            name,
527            DtypeRule::Preserve,
528            GradFlow::Propagates,
529            false,
530            Some("dtype-preserving unary; undefined at 0"),
531        )
532        .with_requires(DomainRequirement::StrictlyPositive),
533        "recip" => OpEffect::known(
534            name,
535            DtypeRule::Preserve,
536            GradFlow::Propagates,
537            false,
538            Some("dtype-preserving unary; undefined at 0"),
539        )
540        .with_requires(DomainRequirement::NonZero),
541        "floor" | "round" | "sign" => OpEffect::known(
542            name,
543            DtypeRule::Preserve,
544            GradFlow::Severs,
545            false,
546            Some("candle operation is created without a backward rule"),
547        ),
548        "ceil" => OpEffect::known(
549            name,
550            DtypeRule::Preserve,
551            GradFlow::Unknown,
552            false,
553            Some("candle backward reports this operation as unsupported"),
554        ),
555        "cmp" | "eq" | "ne" | "lt" | "gt" | "ge" | "le" => OpEffect::known(
556            name,
557            DtypeRule::Fixed(AbstractDtype::U8),
558            GradFlow::Severs,
559            false,
560            Some("candle-core 0.11.0 comparisons return U8 and do not propagate gradients"),
561        ),
562        "argmin" | "argmin_keepdim" | "argmax" | "argmax_keepdim" => OpEffect::known(
563            name,
564            DtypeRule::Fixed(AbstractDtype::U32),
565            GradFlow::Severs,
566            false,
567            Some("candle-core 0.11.0 index reductions return U32 without a backward op"),
568        ),
569
570        // --- dtype-preserving reductions / shape views with known backprop ---
571        "sum" | "sum_keepdim" | "sum_all" | "mean" | "mean_keepdim" | "mean_all" | "max"
572        | "max_keepdim" | "min" | "min_keepdim" | "log_sum_exp" | "powf" | "elu" | "clamp" => {
573            OpEffect::known(
574                name,
575                DtypeRule::Preserve,
576                GradFlow::Propagates,
577                false,
578                Some("dtype-preserving reduction/affine"),
579            )
580        }
581        // Domain is a transfer over (mul, add) literals — see `affine_domain`.
582        "affine" => OpEffect::known(
583            name,
584            DtypeRule::Preserve,
585            GradFlow::Propagates,
586            false,
587            Some("dtype-preserving affine; domain transfers through positive add"),
588        )
589        .with_domain(NumericDomain::Unknown),
590
591        // --- layout / view ops with explicit Candle backward rules ---
592        "reshape" | "flatten_all" | "flatten_to" | "flatten_from" | "squeeze" | "unsqueeze"
593        | "transpose" | "permute" | "narrow" | "contiguous" | "broadcast_as" | "broadcast_left"
594        | "expand" | "t" => OpEffect::known(
595            name,
596            DtypeRule::Preserve,
597            GradFlow::Propagates,
598            false,
599            Some("layout/view operation with an explicit backward rule"),
600        ),
601
602        // --- explicit sever ---
603        "detach" | "as_detached_tensor" => OpEffect::known(
604            name,
605            DtypeRule::Preserve,
606            GradFlow::Severs,
607            false,
608            Some("explicit detach"),
609        ),
610
611        // --- candle-nn losses (expandable bodies via `library_body`, not precomputed verdicts) ---
612        "mse" => OpEffect::known(
613            name,
614            DtypeRule::Preserve,
615            GradFlow::Propagates,
616            true,
617            Some("candle_nn::loss::mse"),
618        ),
619        "nll" => OpEffect::known(
620            name,
621            DtypeRule::Preserve,
622            GradFlow::Propagates,
623            true,
624            Some("candle_nn::loss::nll expects log-probabilities"),
625        ),
626        "cross_entropy" => OpEffect::known(
627            name,
628            DtypeRule::Preserve,
629            GradFlow::Propagates,
630            true,
631            Some("candle_nn::loss::cross_entropy via log_softmax + nll"),
632        ),
633        "huber" => OpEffect::known(
634            name,
635            DtypeRule::Preserve,
636            GradFlow::Propagates,
637            true,
638            Some("candle_nn::loss::huber"),
639        ),
640        "binary_cross_entropy_with_logit" => OpEffect::known(
641            name,
642            DtypeRule::Preserve,
643            GradFlow::Propagates,
644            true,
645            Some(
646                "candle_nn::loss::binary_cross_entropy_with_logit; body expanded when \
647                 Candle 0.11.0 is resolved — see library_body()",
648            ),
649        ),
650
651        // --- candle-nn ops implemented via apply_op*_no_bwd (candle-nn 0.11.0) ---
652        "softmax_last_dim" => OpEffect::known(
653            name,
654            DtypeRule::Preserve,
655            GradFlow::Severs,
656            false,
657            Some("candle_nn::ops::softmax_last_dim uses apply_op1_no_bwd"),
658        )
659        .with_domain(NumericDomain::SaturatingUnit),
660        "rms_norm" => OpEffect::known(
661            name,
662            DtypeRule::Preserve,
663            GradFlow::Severs,
664            false,
665            Some("candle_nn::ops::rms_norm uses apply_op2_no_bwd"),
666        ),
667        "layer_norm" => OpEffect::known(
668            name,
669            DtypeRule::Preserve,
670            GradFlow::Severs,
671            false,
672            Some("candle_nn::ops::layer_norm uses apply_op3_no_bwd"),
673        ),
674        "sdpa" => OpEffect::known(
675            name,
676            DtypeRule::Preserve,
677            GradFlow::Severs,
678            false,
679            Some("candle_nn::ops::sdpa uses apply_op3_no_bwd"),
680        ),
681        "rope" | "rope_i" | "rope_thd" => OpEffect::known(
682            name,
683            DtypeRule::Preserve,
684            GradFlow::Severs,
685            false,
686            Some("candle_nn::rotary_emb::* uses apply_op3_no_bwd"),
687        ),
688        "flash_attn"
689        | "flash_attn_varlen_cpu"
690        | "flash_attn_varlen_unfused"
691        | "run_flash_attn_cpu" => OpEffect::known(
692            name,
693            DtypeRule::Preserve,
694            GradFlow::Severs,
695            false,
696            Some("candle-nn 0.11.0 CPU/varlen attention returns a detached output"),
697        ),
698
699        // --- candle-nn ops that DO have differentiable slow paths ---
700        "rms_norm_slow" | "layer_norm_slow" | "rope_slow" | "rope_i_slow" => OpEffect::known(
701            name,
702            DtypeRule::Preserve,
703            GradFlow::Propagates,
704            false,
705            Some("candle_nn differentiable helper"),
706        ),
707        "sigmoid" => OpEffect::known(
708            name,
709            DtypeRule::Preserve,
710            GradFlow::Propagates,
711            false,
712            Some(
713                "candle_nn sigmoid saturates to exactly 0/1 in f32 \
714                 (candle-nn-0.11.0/src/ops.rs:59-60)",
715            ),
716        )
717        .with_domain(NumericDomain::SaturatingUnit),
718        "softmax" | "log_softmax" => OpEffect::known(
719            name,
720            DtypeRule::Preserve,
721            GradFlow::Propagates,
722            false,
723            Some("candle_nn softmax family; probabilities attain 0/1 after f32 rounding"),
724        )
725        .with_domain(NumericDomain::SaturatingUnit),
726
727        // device moves preserve dtype and grad
728        "to_device" | "clone" => {
729            OpEffect::known(name, DtypeRule::Preserve, GradFlow::Propagates, false, None)
730        }
731
732        _ => OpEffect::unknown(name),
733    }
734}
735
736/// Exact method semantics where the bare method name is insufficient.
737///
738/// In candle-nn 0.11.0 `RmsNorm::forward` chooses a no-backward custom kernel for contiguous
739/// inputs and a differentiable implementation otherwise, whereas `forward_diff` always selects
740/// the differentiable implementation. Treating both as a bare `forward` would lose the crucial
741/// distinction.
742pub fn lookup_method(
743    receiver_type: &str,
744    method: &str,
745    candle_nn_version: Option<&str>,
746) -> OpEffect {
747    let receiver = receiver_type
748        .split('<')
749        .next()
750        .unwrap_or(receiver_type)
751        .rsplit("::")
752        .next()
753        .unwrap_or(receiver_type);
754    match (receiver, method, candle_nn_version) {
755        ("RmsNorm", "forward_diff", Some(AUDITED_CANDLE_VERSION)) => OpEffect::known(
756            "RmsNorm::forward_diff",
757            DtypeRule::Preserve,
758            GradFlow::Propagates,
759            false,
760            Some("audited candle-nn RmsNorm::forward_diff uses differentiable LayerNorm::forward"),
761        ),
762        ("RmsNorm", "forward", Some(AUDITED_CANDLE_VERSION)) => OpEffect::known(
763            "RmsNorm::forward",
764            DtypeRule::Preserve,
765            GradFlow::LayoutDependent,
766            false,
767            Some("audited candle-nn RmsNorm::forward uses a no-bwd kernel for contiguous input"),
768        ),
769        (
770            "Linear" | "Embedding" | "Conv1d" | "Conv2d" | "ConvTranspose1d" | "ConvTranspose2d"
771            | "PReLU" | "GroupNorm" | "BatchNorm",
772            "forward" | "forward_t",
773            Some(AUDITED_CANDLE_VERSION),
774        ) => OpEffect::known(
775            &format!("{receiver}::{method}"),
776            DtypeRule::Preserve,
777            GradFlow::Propagates,
778            false,
779            Some("known candle-nn parameterized module forward"),
780        ),
781        _ => OpEffect::unknown(&format!("{receiver}::{method}")),
782    }
783}
784
785/// Resolve an operation label that may already contain a receiver type.
786///
787/// Dataflow nodes preserve labels such as `RmsNorm::forward`; this helper keeps their exact
788/// method rule instead of collapsing them back to the ambiguous bare name `forward`.
789pub fn lookup_resolved(op: &str, candle_nn_version: Option<&str>) -> OpEffect {
790    if let Some((receiver, method)) = op.rsplit_once("::") {
791        let method_effect = lookup_method(receiver, method, candle_nn_version);
792        if !matches!(method_effect.dtype, DtypeRule::Unknown)
793            || !matches!(method_effect.grad, GradFlow::Unknown)
794        {
795            return method_effect;
796        }
797    }
798    lookup_for(op, candle_nn_version)
799}
800
801/// Whether `op` is a same-dtype binary that should emit a conflict diagnostic on mismatch.
802#[allow(dead_code)]
803pub fn requires_same_dtype(op: &str) -> bool {
804    matches!(lookup(op).dtype, DtypeRule::SameAsInputs)
805}
806
807/// Whether `op` is a known candle-nn helper that severs backward.
808#[allow(dead_code)]
809pub fn is_no_backward(op: &str) -> bool {
810    matches!(lookup(op).grad, GradFlow::Severs)
811}
candle_graph/op_semantics.rs

candle_graph/
op_semantics.rs
