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    /// Dtype known from dataflow on the output but no catalog rule for the callee name.
271    pub fn inferred_preserve(name: &str) -> Self {
272        Self::known(
273            name,
274            DtypeRule::Preserve,
275            GradFlow::Unknown,
276            false,
277            Some("dataflow resolved output dtype; treat as dtype-preserving"),
278        )
279    }
280
281    fn known(
282        name: &str,
283        dtype: DtypeRule,
284        grad: GradFlow,
285        is_loss: bool,
286        note: Option<&'static str>,
287    ) -> Self {
288        Self {
289            name: name.to_string(),
290            dtype,
291            grad,
292            is_loss,
293            note,
294            domain: NumericDomain::Real,
295            requires: DomainRequirement::None,
296        }
297    }
298
299    fn with_domain(mut self, domain: NumericDomain) -> Self {
300        self.domain = domain;
301        self
302    }
303
304    fn with_requires(mut self, requires: DomainRequirement) -> Self {
305        self.requires = requires;
306        self
307    }
308
309    pub fn unknown(name: &str) -> Self {
310        Self {
311            name: name.to_string(),
312            dtype: DtypeRule::Unknown,
313            grad: GradFlow::Unknown,
314            is_loss: false,
315            note: Some("no transfer rule; left Unknown"),
316            domain: NumericDomain::Unknown,
317            requires: DomainRequirement::None,
318        }
319    }
320
321    pub fn domain_rule_label(&self) -> String {
322        match self.domain {
323            NumericDomain::Real => "real".into(),
324            NumericDomain::NonNegative => "non_negative".into(),
325            NumericDomain::SaturatingUnit => format!(
326                "saturating_unit(f32_upper={SIGMOID_F32_UPPER_SATURATION},f32_lower={SIGMOID_F32_LOWER_SATURATION})"
327            ),
328            NumericDomain::StrictlyPositive => "strictly_positive".into(),
329            NumericDomain::Unknown => "unknown".into(),
330        }
331    }
332}
333
334/// Strength of a domain-requirement failure.
335#[derive(Debug, Clone, Copy, PartialEq, Eq)]
336pub enum DomainViolationConfidence {
337    Proven,
338    Unknown,
339}
340
341/// Whether `domain` fails to discharge `requires`.
342///
343/// Returns `Some(Proven)` for a catalog-known violation, `Some(Unknown)` when the producer
344/// domain is unknown, and `None` when the requirement is met.
345pub fn domain_violation(
346    requires: DomainRequirement,
347    domain: NumericDomain,
348) -> Option<DomainViolationConfidence> {
349    match requires {
350        DomainRequirement::None => None,
351        DomainRequirement::StrictlyPositive => match domain {
352            NumericDomain::StrictlyPositive => None,
353            NumericDomain::Unknown => Some(DomainViolationConfidence::Unknown),
354            NumericDomain::Real | NumericDomain::NonNegative | NumericDomain::SaturatingUnit => {
355                Some(DomainViolationConfidence::Proven)
356            }
357        },
358        DomainRequirement::NonZero => match domain {
359            NumericDomain::StrictlyPositive => None,
360            NumericDomain::Unknown => Some(DomainViolationConfidence::Unknown),
361            NumericDomain::Real | NumericDomain::NonNegative | NumericDomain::SaturatingUnit => {
362                Some(DomainViolationConfidence::Proven)
363            }
364        },
365        DomainRequirement::NonNegative => match domain {
366            NumericDomain::NonNegative
367            | NumericDomain::StrictlyPositive
368            | NumericDomain::SaturatingUnit => None,
369            NumericDomain::Unknown => Some(DomainViolationConfidence::Unknown),
370            NumericDomain::Real => Some(DomainViolationConfidence::Proven),
371        },
372    }
373}
374
375/// Join domains through a dtype-preserving unary/binary where both sides contribute.
376pub fn join_domain(left: NumericDomain, right: NumericDomain) -> NumericDomain {
377    use NumericDomain::*;
378    match (left, right) {
379        (Unknown, _) | (_, Unknown) => Unknown,
380        (StrictlyPositive, StrictlyPositive) => StrictlyPositive,
381        (StrictlyPositive, NonNegative) | (NonNegative, StrictlyPositive) => NonNegative,
382        (NonNegative, NonNegative) => NonNegative,
383        (SaturatingUnit, SaturatingUnit) => SaturatingUnit,
384        (SaturatingUnit, NonNegative) | (NonNegative, SaturatingUnit) => NonNegative,
385        (SaturatingUnit, StrictlyPositive) | (StrictlyPositive, SaturatingUnit) => NonNegative,
386        _ => Real,
387    }
388}
389
390/// Transfer for `tensor.affine(mul, add)` with proven literal coefficients.
391///
392/// Only the epsilon-guard form `mul > 0 && add > 0` on a non-negative / unit operand
393/// discharges `StrictlyPositive`. Reflections such as `affine(-1, 1)` (`1 - p`) keep a
394/// zero-attaining domain.
395pub fn affine_domain(operand: NumericDomain, mul: Option<f64>, add: Option<f64>) -> NumericDomain {
396    match (mul, add) {
397        (Some(m), Some(a)) if m > 0.0 && a > 0.0 => match operand {
398            NumericDomain::SaturatingUnit
399            | NumericDomain::NonNegative
400            | NumericDomain::StrictlyPositive => NumericDomain::StrictlyPositive,
401            NumericDomain::Unknown => NumericDomain::Unknown,
402            NumericDomain::Real => NumericDomain::Real,
403        },
404        _ => match operand {
405            // 1 - p over a unit interval still attains 0 and 1.
406            NumericDomain::SaturatingUnit => NumericDomain::SaturatingUnit,
407            NumericDomain::StrictlyPositive => NumericDomain::Real,
408            NumericDomain::NonNegative => NumericDomain::Real,
409            NumericDomain::Unknown => NumericDomain::Unknown,
410            NumericDomain::Real => NumericDomain::Real,
411        },
412    }
413}
414
415pub fn domain_includes_zero(domain: NumericDomain) -> bool {
416    matches!(
417        domain,
418        NumericDomain::NonNegative | NumericDomain::SaturatingUnit | NumericDomain::Real
419    )
420}
421
422/// Last path segment of a call, e.g. `broadcast_add` or `cross_entropy`.
423pub fn lookup(op: &str) -> OpEffect {
424    // Without Cargo evidence the version-sensitive answer is unknown. Callers that resolved the
425    // target crate must use `lookup_for`.
426    lookup_for(op, None)
427}
428
429/// Look up an operation against an audited Candle version.
430///
431/// Generic Tensor algebra is stable across the supported catalog. Rules tied to Candle's
432/// custom `*_no_bwd` implementations are only asserted for versions whose source was audited;
433/// other versions stay `Unknown` rather than inheriting a potentially stale gradient claim.
434pub fn lookup_for(op: &str, candle_nn_version: Option<&str>) -> OpEffect {
435    let name = op.rsplit("::").next().unwrap_or(op);
436    let audited_no_bwd = matches!(
437        name,
438        "softmax_last_dim"
439            | "rms_norm"
440            | "layer_norm"
441            | "sdpa"
442            | "rope"
443            | "rope_i"
444            | "rope_thd"
445            | "flash_attn"
446            | "flash_attn_varlen_cpu"
447            | "flash_attn_varlen_unfused"
448            | "run_flash_attn_cpu"
449    );
450    if audited_no_bwd && !candle_nn_version.is_some_and(is_audited_candle_version) {
451        return OpEffect::unknown(name);
452    }
453    match name {
454        // --- explicit dtype ---
455        "to_dtype" => OpEffect::known(
456            name,
457            DtypeRule::Explicit,
458            GradFlow::Propagates,
459            false,
460            Some("candle Tensor::to_dtype"),
461        ),
462
463        // --- same-dtype binaries (conflict when operands disagree) ---
464        "add" | "sub" | "mul" | "maximum" | "minimum" => OpEffect::known(
465            name,
466            DtypeRule::SameAsInputs,
467            GradFlow::Propagates,
468            false,
469            Some("candle binary op; operands must share dtype"),
470        ),
471        "div" => OpEffect::known(
472            name,
473            DtypeRule::SameAsInputs,
474            GradFlow::Propagates,
475            false,
476            Some("candle binary op; operands must share dtype"),
477        )
478        .with_requires(DomainRequirement::NonZero),
479        "broadcast_add" | "broadcast_sub" | "broadcast_mul" | "broadcast_maximum"
480        | "broadcast_minimum" | "broadcast_pow" => OpEffect::known(
481            name,
482            DtypeRule::SameAsInputs,
483            GradFlow::Propagates,
484            false,
485            Some("candle broadcast binary; operands must share dtype"),
486        ),
487        "broadcast_div" => OpEffect::known(
488            name,
489            DtypeRule::SameAsInputs,
490            GradFlow::Propagates,
491            false,
492            Some("candle broadcast binary; operands must share dtype"),
493        )
494        .with_requires(DomainRequirement::NonZero),
495        "matmul" | "broadcast_matmul" => OpEffect::known(
496            name,
497            DtypeRule::SameAsInputs,
498            GradFlow::Propagates,
499            false,
500            Some("candle matmul; operands must share dtype"),
501        ),
502
503        // --- dtype-preserving unary math ---
504        "neg" | "sin" | "cos" | "tanh" | "gelu" | "silu" | "erf" => OpEffect::known(
505            name,
506            DtypeRule::Preserve,
507            GradFlow::Propagates,
508            false,
509            Some("dtype-preserving unary"),
510        ),
511        "abs" | "sqr" | "relu" => OpEffect::known(
512            name,
513            DtypeRule::Preserve,
514            GradFlow::Propagates,
515            false,
516            Some("dtype-preserving unary"),
517        )
518        .with_domain(NumericDomain::NonNegative),
519        "exp" => OpEffect::known(
520            name,
521            DtypeRule::Preserve,
522            GradFlow::Propagates,
523            false,
524            Some("dtype-preserving unary; f32 underflows to exactly 0.0"),
525        )
526        .with_domain(NumericDomain::NonNegative),
527        "sqrt" => OpEffect::known(
528            name,
529            DtypeRule::Preserve,
530            GradFlow::Propagates,
531            false,
532            Some("dtype-preserving unary"),
533        )
534        .with_domain(NumericDomain::NonNegative)
535        .with_requires(DomainRequirement::NonNegative),
536        "log" => OpEffect::known(
537            name,
538            DtypeRule::Preserve,
539            GradFlow::Propagates,
540            false,
541            Some("dtype-preserving unary; undefined at 0"),
542        )
543        .with_requires(DomainRequirement::StrictlyPositive),
544        "recip" => OpEffect::known(
545            name,
546            DtypeRule::Preserve,
547            GradFlow::Propagates,
548            false,
549            Some("dtype-preserving unary; undefined at 0"),
550        )
551        .with_requires(DomainRequirement::NonZero),
552        "floor" | "round" | "sign" => OpEffect::known(
553            name,
554            DtypeRule::Preserve,
555            GradFlow::Severs,
556            false,
557            Some("candle operation is created without a backward rule"),
558        ),
559        "ceil" => OpEffect::known(
560            name,
561            DtypeRule::Preserve,
562            GradFlow::Unknown,
563            false,
564            Some("candle backward reports this operation as unsupported"),
565        ),
566        "cmp" | "eq" | "ne" | "lt" | "gt" | "ge" | "le" => OpEffect::known(
567            name,
568            DtypeRule::Fixed(AbstractDtype::U8),
569            GradFlow::Severs,
570            false,
571            Some("candle-core 0.11.0 comparisons return U8 and do not propagate gradients"),
572        ),
573        "argmin" | "argmin_keepdim" | "argmax" | "argmax_keepdim" => OpEffect::known(
574            name,
575            DtypeRule::Fixed(AbstractDtype::U32),
576            GradFlow::Severs,
577            false,
578            Some("candle-core 0.11.0 index reductions return U32 without a backward op"),
579        ),
580
581        // --- dtype-preserving reductions / shape views with known backprop ---
582        "sum" | "sum_keepdim" | "sum_all" | "mean" | "mean_keepdim" | "mean_all" | "max"
583        | "max_keepdim" | "min" | "min_keepdim" | "log_sum_exp" | "powf" | "elu" | "clamp" => {
584            OpEffect::known(
585                name,
586                DtypeRule::Preserve,
587                GradFlow::Propagates,
588                false,
589                Some("dtype-preserving reduction/affine"),
590            )
591        }
592        // Domain is a transfer over (mul, add) literals — see `affine_domain`.
593        "affine" => OpEffect::known(
594            name,
595            DtypeRule::Preserve,
596            GradFlow::Propagates,
597            false,
598            Some("dtype-preserving affine; domain transfers through positive add"),
599        )
600        .with_domain(NumericDomain::Unknown),
601
602        // --- layout / view ops with explicit Candle backward rules ---
603        "reshape" | "flatten_all" | "flatten_to" | "flatten_from" | "squeeze" | "unsqueeze"
604        | "transpose" | "permute" | "narrow" | "contiguous" | "broadcast_as" | "broadcast_left"
605        | "expand" | "t" => OpEffect::known(
606            name,
607            DtypeRule::Preserve,
608            GradFlow::Propagates,
609            false,
610            Some("layout/view operation with an explicit backward rule"),
611        ),
612
613        // --- explicit sever ---
614        "detach" | "as_detached_tensor" => OpEffect::known(
615            name,
616            DtypeRule::Preserve,
617            GradFlow::Severs,
618            false,
619            Some("explicit detach"),
620        ),
621
622        // --- candle-nn losses (expandable bodies via `library_body`, not precomputed verdicts) ---
623        "mse" => OpEffect::known(
624            name,
625            DtypeRule::Preserve,
626            GradFlow::Propagates,
627            true,
628            Some("candle_nn::loss::mse"),
629        ),
630        "nll" => OpEffect::known(
631            name,
632            DtypeRule::Preserve,
633            GradFlow::Propagates,
634            true,
635            Some("candle_nn::loss::nll expects log-probabilities"),
636        ),
637        "cross_entropy" => OpEffect::known(
638            name,
639            DtypeRule::Preserve,
640            GradFlow::Propagates,
641            true,
642            Some("candle_nn::loss::cross_entropy via log_softmax + nll"),
643        ),
644        "huber" => OpEffect::known(
645            name,
646            DtypeRule::Preserve,
647            GradFlow::Propagates,
648            true,
649            Some("candle_nn::loss::huber"),
650        ),
651        "binary_cross_entropy_with_logit" => OpEffect::known(
652            name,
653            DtypeRule::Preserve,
654            GradFlow::Propagates,
655            true,
656            Some(
657                "candle_nn::loss::binary_cross_entropy_with_logit; body expanded when \
658                 Candle 0.11.0 is resolved — see library_body()",
659            ),
660        ),
661
662        // --- candle-nn ops implemented via apply_op*_no_bwd (candle-nn 0.11.0) ---
663        "softmax_last_dim" => OpEffect::known(
664            name,
665            DtypeRule::Preserve,
666            GradFlow::Severs,
667            false,
668            Some("candle_nn::ops::softmax_last_dim uses apply_op1_no_bwd"),
669        )
670        .with_domain(NumericDomain::SaturatingUnit),
671        "rms_norm" => OpEffect::known(
672            name,
673            DtypeRule::Preserve,
674            GradFlow::Severs,
675            false,
676            Some("candle_nn::ops::rms_norm uses apply_op2_no_bwd"),
677        ),
678        "layer_norm" => OpEffect::known(
679            name,
680            DtypeRule::Preserve,
681            GradFlow::Severs,
682            false,
683            Some("candle_nn::ops::layer_norm uses apply_op3_no_bwd"),
684        ),
685        "sdpa" => OpEffect::known(
686            name,
687            DtypeRule::Preserve,
688            GradFlow::Severs,
689            false,
690            Some("candle_nn::ops::sdpa uses apply_op3_no_bwd"),
691        ),
692        "rope" | "rope_i" | "rope_thd" => OpEffect::known(
693            name,
694            DtypeRule::Preserve,
695            GradFlow::Severs,
696            false,
697            Some("candle_nn::rotary_emb::* uses apply_op3_no_bwd"),
698        ),
699        "flash_attn"
700        | "flash_attn_varlen_cpu"
701        | "flash_attn_varlen_unfused"
702        | "run_flash_attn_cpu" => OpEffect::known(
703            name,
704            DtypeRule::Preserve,
705            GradFlow::Severs,
706            false,
707            Some("candle-nn 0.11.0 CPU/varlen attention returns a detached output"),
708        ),
709
710        // --- candle-nn ops that DO have differentiable slow paths ---
711        "rms_norm_slow" | "layer_norm_slow" | "rope_slow" | "rope_i_slow" => OpEffect::known(
712            name,
713            DtypeRule::Preserve,
714            GradFlow::Propagates,
715            false,
716            Some("candle_nn differentiable helper"),
717        ),
718        "sigmoid" => OpEffect::known(
719            name,
720            DtypeRule::Preserve,
721            GradFlow::Propagates,
722            false,
723            Some(
724                "candle_nn sigmoid saturates to exactly 0/1 in f32 \
725                 (candle-nn-0.11.0/src/ops.rs:59-60)",
726            ),
727        )
728        .with_domain(NumericDomain::SaturatingUnit),
729        "softmax" | "log_softmax" => OpEffect::known(
730            name,
731            DtypeRule::Preserve,
732            GradFlow::Propagates,
733            false,
734            Some("candle_nn softmax family; probabilities attain 0/1 after f32 rounding"),
735        )
736        .with_domain(NumericDomain::SaturatingUnit),
737
738        // Result/type wrappers and views preserve tensor metadata.
739        "map_err" | "ok" | "err" => OpEffect::known(
740            name,
741            DtypeRule::Preserve,
742            GradFlow::Propagates,
743            false,
744            Some("Result/type wrapper preserves tensor dtype"),
745        ),
746
747        // device moves preserve dtype and grad
748        "to_device" | "clone" => {
749            OpEffect::known(name, DtypeRule::Preserve, GradFlow::Propagates, false, None)
750        }
751
752        _ => OpEffect::unknown(name),
753    }
754}
755
756/// Exact method semantics where the bare method name is insufficient.
757///
758/// In candle-nn 0.11.0 `RmsNorm::forward` chooses a no-backward custom kernel for contiguous
759/// inputs and a differentiable implementation otherwise, whereas `forward_diff` always selects
760/// the differentiable implementation. Treating both as a bare `forward` would lose the crucial
761/// distinction.
762pub fn lookup_method(
763    receiver_type: &str,
764    method: &str,
765    candle_nn_version: Option<&str>,
766) -> OpEffect {
767    let receiver = receiver_type
768        .split('<')
769        .next()
770        .unwrap_or(receiver_type)
771        .rsplit("::")
772        .next()
773        .unwrap_or(receiver_type);
774    match (receiver, method, candle_nn_version) {
775        ("RmsNorm", "forward_diff", Some(AUDITED_CANDLE_VERSION)) => OpEffect::known(
776            "RmsNorm::forward_diff",
777            DtypeRule::Preserve,
778            GradFlow::Propagates,
779            false,
780            Some("audited candle-nn RmsNorm::forward_diff uses differentiable LayerNorm::forward"),
781        ),
782        ("RmsNorm", "forward", Some(AUDITED_CANDLE_VERSION)) => OpEffect::known(
783            "RmsNorm::forward",
784            DtypeRule::Preserve,
785            GradFlow::LayoutDependent,
786            false,
787            Some("audited candle-nn RmsNorm::forward uses a no-bwd kernel for contiguous input"),
788        ),
789        (
790            "Linear" | "Embedding" | "Conv1d" | "Conv2d" | "ConvTranspose1d" | "ConvTranspose2d"
791            | "PReLU" | "GroupNorm" | "BatchNorm",
792            "forward" | "forward_t",
793            Some(AUDITED_CANDLE_VERSION),
794        ) => OpEffect::known(
795            &format!("{receiver}::{method}"),
796            DtypeRule::Preserve,
797            GradFlow::Propagates,
798            false,
799            Some("known candle-nn parameterized module forward"),
800        ),
801        _ => OpEffect::unknown(&format!("{receiver}::{method}")),
802    }
803}
804
805/// Resolve an operation label that may already contain a receiver type.
806///
807/// Dataflow nodes preserve labels such as `RmsNorm::forward`; this helper keeps their exact
808/// method rule instead of collapsing them back to the ambiguous bare name `forward`.
809pub fn lookup_resolved(op: &str, candle_nn_version: Option<&str>) -> OpEffect {
810    if let Some((receiver, method)) = op.rsplit_once("::") {
811        let method_effect = lookup_method(receiver, method, candle_nn_version);
812        if !matches!(method_effect.dtype, DtypeRule::Unknown)
813            || !matches!(method_effect.grad, GradFlow::Unknown)
814        {
815            return method_effect;
816        }
817    }
818    lookup_for(op, candle_nn_version)
819}
820
821/// Whether `op` is a same-dtype binary that should emit a conflict diagnostic on mismatch.
822#[allow(dead_code)]
823pub fn requires_same_dtype(op: &str) -> bool {
824    matches!(lookup(op).dtype, DtypeRule::SameAsInputs)
825}
826
827/// Whether `op` is a known candle-nn helper that severs backward.
828#[allow(dead_code)]
829pub fn is_no_backward(op: &str) -> bool {
830    matches!(lookup(op).grad, GradFlow::Severs)
831}
candle_graph/op_semantics.rs

candle_graph/
op_semantics.rs
