tract_core/ops/
logic.rs

1#![allow(clippy::bool_comparison)]
2#![allow(clippy::unnecessary_cast)]
3
4mod comparison;
5mod ite;
6pub use comparison::{CompEq, CompGT, CompGTE, CompLT, CompLTE, CompNE};
7pub use comparison::{comp_eq, comp_gt, comp_gte, comp_lt, comp_lte, comp_ne};
8pub use ite::IfThenElse;
9
10use ndarray::*;
11
12use crate::broadcast::multi_broadcast;
13use crate::internal::*;
14
15bin_to_super_type!(and, And,
16                   neutral_element: 1,
17                   absorbing_element: 0,
18                   [bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = (a as i64 != 0 && b as i64 != 0) as _);
19bin_to_super_type!(or, Or,
20                   neutral_element: 0,
21                   absorbing_element: 1,
22                   [bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = (a as i64 != 0 || b as i64 != 0) as _);
23bin_to_super_type!(xor, Xor, declutter: declutter_xor, neutral_element: 0, [bool] => |c, &a, &b| *c = a ^ b);
24
25fn declutter_xor(
26    _op: &Xor,
27    model: &TypedModel,
28    node: &TypedNode,
29) -> TractResult<Option<TypedModelPatch>> {
30    // Xor(x, 1) = Not(x)
31    if let Some(uniform) = crate::ops::binary::one_input_is_uniform(model, node)?
32        && tensor0(1i64).close_enough(&uniform.uni, false).is_ok()
33    {
34        return Ok(Some(TypedModelPatch::replace_single_op(
35            model,
36            node,
37            &[uniform.var],
38            crate::ops::element_wise::ElementWiseOp(Box::new(Not {}), None),
39        )?));
40    }
41    Ok(None)
42}
43
44element_wise!(not, Not, [bool] => |_, vs| {
45    vs.iter_mut().for_each(|a| *a = !*a);
46    Ok(())
47});
48
49#[derive(Debug, Clone, new, Default, Hash, PartialEq, Eq)]
50pub struct Iff;
51
52impl Iff {
53    pub unsafe fn eval_t<T: Datum>(
54        cond: &ArrayViewD<bool>,
55        out: &mut Tensor,
56        t: &Tensor,
57        f: &Tensor,
58    ) {
59        unsafe {
60            Zip::from(out.to_array_view_mut_unchecked::<T>())
61                .and_broadcast(cond)
62                .and_broadcast(t.to_array_view_unchecked::<T>())
63                .and_broadcast(f.to_array_view_unchecked::<T>())
64                .for_each(|r, c, t, f| *r = if *c { t.clone() } else { f.clone() })
65        }
66    }
67}
68
69impl Op for Iff {
70    fn name(&self) -> StaticName {
71        "Iff".into()
72    }
73    op_as_typed_op!();
74}
75
76impl EvalOp for Iff {
77    fn is_stateless(&self) -> bool {
78        true
79    }
80
81    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
82        let (cond, t, f) = args_3!(inputs);
83        anyhow::ensure!(t.datum_type() == f.datum_type());
84        let shape: TVec<usize> = multi_broadcast(&[cond.shape(), t.shape(), f.shape()])?;
85        unsafe {
86            let mut result = Tensor::uninitialized_dt(t.datum_type(), &shape)?;
87            let cond = cond.to_plain_array_view::<bool>()?;
88            dispatch_datum_by_size!(Self::eval_t(t.datum_type())(&cond, &mut result, &t, &f));
89            Ok(tvec!(result.into_tvalue()))
90        }
91    }
92}
93
94pub fn sym_to_coord_axis(sym: &Symbol) -> Option<usize> {
95    format!("{sym}").strip_prefix("🎯")?.parse::<usize>().ok()
96}
97
98pub(crate) fn coord_bound_assertions(expr: &TDim, shape: &ShapeFact) -> Vec<Assertion> {
99    expr.symbols()
100        .into_iter()
101        .filter_map(|s| sym_to_coord_axis(&s).filter(|k| *k < shape.rank()).map(|k| (k, s)))
102        .flat_map(|(k, sym)| {
103            [
104                Assertion::GTE(TDim::Sym(sym.clone()), TDim::Val(0)),
105                Assertion::LTE(TDim::Sym(sym), shape[k].clone() - TDim::Val(1)),
106            ]
107        })
108        .collect()
109}
110
111pub(crate) fn is_provably_all_false(expr: &TDim, shape: &ShapeFact) -> bool {
112    let extra = coord_bound_assertions(expr, shape);
113    expr.clone().simplify_with_extra_assertions(&extra) == TDim::Val(0)
114}
115
116pub(crate) fn is_provably_all_true(expr: &TDim, shape: &ShapeFact) -> bool {
117    let extra = coord_bound_assertions(expr, shape);
118    expr.clone().simplify_with_extra_assertions(&extra) == TDim::Val(1)
119}
120
121/// The interval of indices along one axis where a boolean condition is true.
122///
123/// `None` on a bound means "open" — start defaults to 0, end defaults to `dim`.
124///
125/// | `start`   | `end`        | meaning                           |
126/// |-----------|--------------|-----------------------------------|
127/// | `None`    | `None`       | whole dimension (AllTrue)         |
128/// | `None`    | `Some(0)`    | empty (AllFalse)                  |
129/// | `None`    | `Some(e)`    | `[0, e)` — lower region true      |
130/// | `Some(s)` | `None`       | `[s, dim)` — upper region true    |
131/// | `Some(s)` | `Some(e)`    | `[s, e)` — three zones            |
132#[derive(Debug, Clone)]
133pub(crate) struct TrueRange {
134    pub axis: usize,
135    pub start: Option<TDim>, // None = 0
136    pub end: Option<TDim>,   // None = dim
137}
138
139impl TrueRange {
140    /// Condition is true for the entire dimension.
141    pub fn is_full(&self) -> bool {
142        self.start.is_none() && self.end.is_none()
143    }
144    /// Condition is never true (empty range).
145    pub fn is_empty(&self) -> bool {
146        match (&self.start, &self.end) {
147            (None, Some(e)) => *e == TDim::Val(0),
148            (Some(s), Some(e)) => s == e,
149            _ => false,
150        }
151    }
152}
153
154pub(crate) fn classify_true_range(expr: &TDim, shape: &ShapeFact) -> Option<TrueRange> {
155    fn try_ge(ge: &TDim, shape: &ShapeFact) -> Option<(usize, TDim)> {
156        if let TDim::Ge(lhs, rhs) = ge
157            && let TDim::Sym(sym) = &**lhs
158        {
159            let k = sym_to_coord_axis(sym)?;
160            if k < shape.rank() && !rhs.symbols().contains(sym) {
161                return Some((k, *rhs.clone()));
162            }
163        }
164        None
165    }
166
167    let simplified = expr.clone().simplify();
168    // All-false: empty range on axis 0
169    if simplified == TDim::Val(0) || is_provably_all_false(&simplified, shape) {
170        return Some(TrueRange { axis: 0, start: None, end: Some(TDim::Val(0)) });
171    }
172    // All-true: open (unbounded) range on axis 0
173    if simplified == TDim::Val(1) || is_provably_all_true(&simplified, shape) {
174        return Some(TrueRange { axis: 0, start: None, end: None });
175    }
176    // Ge(x_k, split): true when x_k >= split → [split, dim)
177    if let Some((axis, split)) = try_ge(&simplified, shape) {
178        return Some(TrueRange { axis, start: Some(split), end: None });
179    }
180    // 1 - Ge(x_k, split): true when x_k < split → [0, split)
181    let flipped = (TDim::Val(1) - simplified).simplify();
182    if let Some((axis, split)) = try_ge(&flipped, shape) {
183        return Some(TrueRange { axis, start: None, end: Some(split) });
184    }
185    None
186}
187
188impl TypedOp for Iff {
189    as_op!();
190
191    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
192        ensure!(inputs.len() == 3, "Iff expects 3 intputs.");
193        ensure!(inputs[1].datum_type == inputs[2].datum_type);
194        ensure!(inputs[0].datum_type.is::<bool>());
195        ensure!(inputs[0].rank() == inputs[1].rank());
196        ensure!(inputs[0].rank() == inputs[2].rank());
197        let shape = multi_broadcast(&[
198            inputs[0].shape.to_tvec(),
199            inputs[1].shape.to_tvec(),
200            inputs[2].shape.to_tvec(),
201        ])
202        .unwrap();
203        let mut fact = inputs[1].datum_type.fact(shape);
204        // Propagate uniform_tdim when condition is provably constant
205        fact.uniform_tdim = match inputs[0].uniform_tdim.as_ref().map(|d| d.clone().simplify()) {
206            Some(TDim::Val(0)) => inputs[2].uniform_tdim.clone(), // always false → false branch
207            Some(TDim::Val(_)) => inputs[1].uniform_tdim.clone(), // always true → true branch
208            _ => None,
209        };
210        Ok(tvec!(fact))
211    }
212
213    fn input_roi(
214        &self,
215        model: &TypedModel,
216        node: &TypedNode,
217    ) -> TractResult<Option<TVec<Option<TDim>>>> {
218        // select(cond, then, else):
219        //   then-branch matters where cond is nonzero → propagate cond
220        //   else-branch matters where cond is zero    → propagate cond==0
221        let cond_fact = model.outlet_fact(node.inputs[0])?;
222        if let Some(cond_expr) = &cond_fact.uniform_tdim {
223            let cond = cond_expr.clone().simplify();
224            let not_cond = TDim::Eq(Box::new(cond.clone()), Box::new(TDim::Val(0))).simplify();
225            return Ok(Some(tvec![None, Some(cond), Some(not_cond)]));
226        }
227        // Bubbling: delegate to the natural blanket implementation.
228        crate::optim::propagate_roi::bubble_roi(model, node)
229    }
230
231    fn declutter(
232        &self,
233        model: &TypedModel,
234        node: &TypedNode,
235    ) -> TractResult<Option<TypedModelPatch>> {
236        // Fold Iff(const, t, f) → t or f.
237        // Symbolic uniform_tdim cases are handled upstream by FoldUniformMask,
238        // which injects a concrete Const(0/1) that this rule then folds.
239        let cond_fact = model.outlet_fact(node.inputs[0])?;
240        rule_if_some!(uniform = &cond_fact.uniform);
241        rule_if_let!(Ok(cond_val) = uniform.cast_to_scalar::<bool>());
242        let branch = if cond_val { node.inputs[1] } else { node.inputs[2] };
243        let mut patch = TypedModelPatch::default();
244        let mut wire = patch.tap_model(model, branch)?;
245        // The output shape is the broadcast of all three inputs; the selected
246        // branch may be narrower and must be broadcast up, not shunted as is.
247        let out_shape = &model.outlet_fact(node.id.into())?.shape;
248        if &model.outlet_fact(branch)?.shape != out_shape {
249            wire = patch.wire_node(
250                format!("{}.broadcast", node.name),
251                crate::ops::array::MultiBroadcastTo::new(out_shape.clone()),
252                &[wire],
253            )?[0];
254        }
255        patch.shunt_outside(model, node.id.into(), wire)?;
256        Ok(Some(patch))
257    }
258
259    fn axes_mapping(
260        &self,
261        inputs: &[&TypedFact],
262        outputs: &[&TypedFact],
263    ) -> TractResult<AxesMapping> {
264        AxesMapping::natural(inputs, outputs)
265    }
266}
267
268bin_to_super_type!(bitand, BitAnd,
269                   absorbing_element: 0,
270                   [bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = a & b);
271bin_to_super_type!(bitor, BitOr,
272                   neutral_element: 0,
273                   [bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = a | b);
274bin_to_super_type!(bitxor, BitXor,
275                   declutter: declutter_bitxor,
276                   neutral_element: 0,
277                   [bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = a ^ b);
278
279fn declutter_bitxor(
280    _op: &BitXor,
281    model: &TypedModel,
282    node: &TypedNode,
283) -> TractResult<Option<TypedModelPatch>> {
284    // BitXor(x, all_ones) = BitNot(x) — for bool, all_ones = 1
285    if let Some(uniform) = crate::ops::binary::one_input_is_uniform(model, node)? {
286        let var_dt = model.outlet_fact(uniform.var)?.datum_type;
287        let is_all_ones = if var_dt.is::<bool>() {
288            tensor0(1i64).close_enough(&uniform.uni, false).is_ok()
289        } else {
290            tensor0(-1i64).close_enough(&uniform.uni, false).is_ok()
291        };
292        if is_all_ones {
293            return Ok(Some(TypedModelPatch::replace_single_op(
294                model,
295                node,
296                &[uniform.var],
297                crate::ops::element_wise::ElementWiseOp(Box::new(BitNot {}), None),
298            )?));
299        }
300    }
301    Ok(None)
302}
303
304element_wise!(bitnot, BitNot, [bool, u8, u16, u32, u64, i8, i16, i32, i64] => |_, xs| {
305    xs.iter_mut().for_each(|x| *x = !*x);
306    Ok(())
307});
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312    use crate::ops::array::TypedConcat;
313    use crate::ops::binary::TypedBinOp;
314    use crate::ops::change_axes::AxisOp;
315
316    /// Test Case 1: Iff where condition is Eq(T, 0) with T >= 1 assertion.
317    /// After declutter, the Iff should fold to the false branch (inputs[2]).
318    #[test]
319    fn iff_fold_case1_eq_t_zero() -> TractResult<()> {
320        let mut model = TypedModel::default();
321        model.symbols.add_assertion("T >= 1")?;
322        let t_sym = model.symbols.sym("T");
323        let t_dim = TDim::Sym(t_sym.clone());
324
325        // Const T (scalar TDim)
326        let t_wire = model.wire_node(
327            "T",
328            crate::ops::konst::Const::new(tensor0(t_dim.clone()).into_arc_tensor())?,
329            &[],
330        )?[0];
331
332        // Const 0 (scalar TDim)
333        let zero_wire = model.wire_node(
334            "zero",
335            crate::ops::konst::Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?,
336            &[],
337        )?[0];
338
339        // Eq(T, 0) → bool scalar
340        let eq_wire = model.wire_node("eq", TypedBinOp(comp_eq(), None), &[t_wire, zero_wire])?[0];
341
342        // Some data wire for the false branch
343        let data_wire = model.add_source("data", TDim::datum_type().scalar_fact())?;
344
345        // Iff(eq, zero, data) — zero is "true" branch, data is "false" branch
346        let iff_wire = model.wire_node("iff", Iff, &[eq_wire, zero_wire, data_wire])?[0];
347        model.select_output_outlets(&[iff_wire])?;
348
349        let model = model.into_decluttered()?;
350
351        // The Iff should have been folded away (condition is always false given T >= 1)
352        let iff_count = model.nodes().iter().filter(|n| n.op_as::<Iff>().is_some()).count();
353        assert_eq!(iff_count, 0, "Expected Iff to be folded, but found {iff_count} Iff nodes");
354        Ok(())
355    }
356
357    /// Test Case 2: range(0,T,1) → unsqueeze(0) → lt(_, T_unsqueezed) → bitnot → Iff
358    /// The bitnot produces Ge(x1, T), all-false for x1 in [0, T-1].
359    /// After declutter, the Iff should fold to the false branch (data input).
360    #[test]
361    fn iff_fold_case2_not_lt_x1_t() -> TractResult<()> {
362        use crate::ops::array::Range;
363
364        let mut model = TypedModel::default();
365        model.symbols.add_assertion("T >= 1")?;
366        let t_sym = model.symbols.sym("T");
367        let t_dim = TDim::Sym(t_sym.clone());
368
369        // Const start=0 (TDim) and step=1 (TDim) — these get uniform_tdim set in output_facts
370        let start = model.wire_node(
371            "start",
372            crate::ops::konst::Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?,
373            &[],
374        )?[0];
375        let step = model.wire_node(
376            "step",
377            crate::ops::konst::Const::new(tensor0(TDim::Val(1)).into_arc_tensor())?,
378            &[],
379        )?[0];
380        // T is a dynamic TDim input (not a Const) so Range takes the else branch and
381        // sets uniform_tdim = start + step * x0 = x0
382        let end = model.add_source("T_dyn", TDim::datum_type().scalar_fact())?;
383
384        // Range(start=0, end=T, step=1) → [T] TDim with uniform_tdim = x0
385        let range = model.wire_node("range", Range::new(t_dim.clone()), &[start, end, step])?[0];
386
387        // unsqueeze(0) → [1, T] TDim, remap x0→x1 → uniform_tdim = x1
388        let range_unsq = model.wire_node("range_unsq", AxisOp::Add(0), &[range])?[0];
389
390        // T const for comparison, scalar TDim with uniform_tdim = Sym(T)
391        let t_const = model.wire_node(
392            "T_const",
393            crate::ops::konst::Const::new(tensor0(t_dim.clone()).into_arc_tensor())?,
394            &[],
395        )?[0];
396        // unsqueeze T_const → [1,1] TDim to match range_unsq rank
397        let t_unsq = model.wire_node("T_unsq", AxisOp::Add(0), &[t_const])?[0];
398        let t_unsq2 = model.wire_node("T_unsq2", AxisOp::Add(0), &[t_unsq])?[0];
399
400        // lt(range_unsq=[1,T], t_unsq2=[1,1]) → bool [1,T], uniform_tdim = Lt(x1,T)
401        let lt = model.wire_node("lt", TypedBinOp(comp_lt(), None), &[range_unsq, t_unsq2])?[0];
402
403        // bitnot(lt): BitNot doesn't propagate uniform_tdim in output_facts,
404        // but Iff::declutter traces through it to get Not(Lt(x1,T))=Ge(x1,T)
405        let bn = model.wire_node("bitnot", bitnot(), &[lt])?[0];
406
407        // Data source [1, T]
408        let data_shape = tvec![TDim::Val(1), t_dim.clone()];
409        let data = model.add_source("data", TDim::datum_type().fact(data_shape.clone()))?;
410
411        // zeros broadcast to [1, T], uniform_tdim = Val(0)
412        let zero_scalar = model.wire_node(
413            "zero_s",
414            crate::ops::konst::Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?,
415            &[],
416        )?[0];
417        let zeros = model.wire_node(
418            "zeros",
419            crate::ops::array::MultiBroadcastTo {
420                shape: ShapeFact::from_dims(data_shape.iter().cloned()),
421            },
422            &[zero_scalar],
423        )?[0];
424
425        // Iff(bn, zeros, data): condition Ge(x1,T) is all-false → fold to data
426        let iff = model.wire_node("iff", Iff, &[bn, zeros, data])?[0];
427        model.select_output_outlets(&[iff])?;
428
429        let model = model.into_decluttered()?;
430
431        let iff_count = model.nodes().iter().filter(|n| n.op_as::<Iff>().is_some()).count();
432        assert_eq!(iff_count, 0, "Expected Iff to be folded, but found {iff_count} Iff nodes");
433        Ok(())
434    }
435
436    /// Rule 2: condition ge(x2, T/160) over [1,1,1+T/160] → slice+concat, no Iff remaining.
437    #[test]
438    fn iff_split_to_slice_concat() -> TractResult<()> {
439        use crate::ops::array::Range;
440
441        let mut model = TypedModel::default();
442        model.symbols.add_assertion("T >= 160")?;
443        let t_sym = model.symbols.sym("T");
444        let t_dim = TDim::Sym(t_sym.clone());
445
446        // split = T/160
447        let split = t_dim.clone() / 160;
448        // output shape: [1, 1, 1 + T/160]
449        let out_len = TDim::Val(1) + split.clone();
450
451        // Build condition: Range over [0, 1+T/160) on axis 2, then compare >= T/160
452        // We'll construct it directly as a source with the right uniform_tdim.
453        // Simpler: use Range + unsqueeze twice + Ge comparison.
454
455        // Range(0, 1+T/160, 1) → [1+T/160] with uniform_tdim = x0
456        let start = model.wire_node(
457            "start",
458            crate::ops::konst::Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?,
459            &[],
460        )?[0];
461        let step = model.wire_node(
462            "step",
463            crate::ops::konst::Const::new(tensor0(TDim::Val(1)).into_arc_tensor())?,
464            &[],
465        )?[0];
466        let end_val = model.wire_node(
467            "end_val",
468            crate::ops::konst::Const::new(tensor0(out_len.clone()).into_arc_tensor())?,
469            &[],
470        )?[0];
471        let range =
472            model.wire_node("range", Range::new(out_len.clone()), &[start, end_val, step])?[0];
473        // unsqueeze(0): [1, 1+T/160], x0 → x1
474        let r1 = model.wire_node("r1", AxisOp::Add(0), &[range])?[0];
475        // unsqueeze(0): [1, 1, 1+T/160], x1 → x2
476        let r2 = model.wire_node("r2", AxisOp::Add(0), &[r1])?[0];
477
478        // split const
479        let split_const = model.wire_node(
480            "split_const",
481            crate::ops::konst::Const::new(tensor0(split.clone()).into_arc_tensor())?,
482            &[],
483        )?[0];
484        // unsqueeze three times so it can broadcast against [1,1,1+T/160]
485        let sc1 = model.wire_node("sc1", AxisOp::Add(0), &[split_const])?[0];
486        let sc2 = model.wire_node("sc2", AxisOp::Add(0), &[sc1])?[0];
487        let sc2 = model.wire_node("sc3", AxisOp::Add(0), &[sc2])?[0];
488
489        // Ge(range_3d, split_3d) → bool [1,1,1+T/160], uniform_tdim = Ge(x2, T/160)
490        let cond = model.wire_node("cond", TypedBinOp(comp_gte(), None), &[r2, sc2])?[0];
491
492        // true and false branches: shape [1,1,1+T/160]
493        let true_branch = model.add_source(
494            "true_b",
495            TDim::datum_type().fact(tvec![TDim::Val(1), TDim::Val(1), out_len.clone()]),
496        )?;
497        let false_branch = model.add_source(
498            "false_b",
499            TDim::datum_type().fact(tvec![TDim::Val(1), TDim::Val(1), out_len.clone()]),
500        )?;
501
502        let iff = model.wire_node("iff", Iff, &[cond, true_branch, false_branch])?[0];
503        model.select_output_outlets(&[iff])?;
504
505        let model = model.into_decluttered()?;
506
507        let iff_count = model.nodes().iter().filter(|n| n.op_as::<Iff>().is_some()).count();
508        assert_eq!(iff_count, 0, "Expected no Iff nodes after declutter, found {iff_count}");
509
510        let concat_count =
511            model.nodes().iter().filter(|n| n.op_as::<TypedConcat>().is_some()).count();
512        assert!(concat_count > 0, "Expected at least one Concat node after declutter");
513
514        Ok(())
515    }
516
517    /// Verify that uniform_tdim propagation produces the expected values at each stage.
518    #[test]
519    fn verify_uniform_tdim_propagation() -> TractResult<()> {
520        use crate::ops::array::Range;
521
522        let mut model = TypedModel::default();
523        model.symbols.add_assertion("T >= 1")?;
524        let t_sym = model.symbols.sym("T");
525        let t_dim = TDim::Sym(t_sym.clone());
526
527        let start = model.wire_node(
528            "start",
529            crate::ops::konst::Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?,
530            &[],
531        )?[0];
532        let step = model.wire_node(
533            "step",
534            crate::ops::konst::Const::new(tensor0(TDim::Val(1)).into_arc_tensor())?,
535            &[],
536        )?[0];
537        let end = model.add_source("T_dyn", TDim::datum_type().scalar_fact())?;
538        let range = model.wire_node("range", Range::new(t_dim.clone()), &[start, end, step])?[0];
539        let range_unsq = model.wire_node("range_unsq", AxisOp::Add(0), &[range])?[0];
540        let t_const = model.wire_node(
541            "T_const",
542            crate::ops::konst::Const::new(tensor0(t_dim.clone()).into_arc_tensor())?,
543            &[],
544        )?[0];
545        let t_unsq = model.wire_node("T_unsq", AxisOp::Add(0), &[t_const])?[0];
546        let t_unsq2 = model.wire_node("T_unsq2", AxisOp::Add(0), &[t_unsq])?[0];
547        let lt = model.wire_node("lt", TypedBinOp(comp_lt(), None), &[range_unsq, t_unsq2])?[0];
548
549        let range_fact = model.outlet_fact(range)?;
550        let range_unsq_fact = model.outlet_fact(range_unsq)?;
551        let t_unsq_fact = model.outlet_fact(t_unsq)?;
552        let lt_fact = model.outlet_fact(lt)?;
553
554        assert!(range_fact.uniform_tdim.is_some(), "range should have uniform_tdim");
555        assert!(range_unsq_fact.uniform_tdim.is_some(), "range_unsq should have uniform_tdim");
556        assert!(t_unsq_fact.uniform_tdim.is_some(), "t_unsq should have uniform_tdim");
557        assert!(lt_fact.uniform_tdim.is_some(), "lt should have uniform_tdim");
558
559        Ok(())
560    }
561
562    /// Iff(const_cond, t, f) where the selected branch is narrower than the
563    /// output (which is the broadcast of cond and both branches).  The fold
564    /// must broadcast the branch up; shunting it as-is fails patch
565    /// validation ("Trying to substitute a 1,2,3 by 1,2,1").
566    #[test]
567    fn iff_fold_broadcasts_narrower_branch() -> TractResult<()> {
568        let mut model = TypedModel::default();
569        let cond = model.wire_node(
570            "cond",
571            crate::ops::konst::Const::new(
572                Tensor::from_shape(&[1, 2, 3], &[false; 6])?.into_arc_tensor(),
573            )?,
574            &[],
575        )?[0];
576        let then = model.add_source("then", f32::fact([1, 2, 3]))?;
577        let otherwise = model.add_source("else", f32::fact([1, 2, 1]))?;
578        let iff = model.wire_node("iff", Iff, &[cond, then, otherwise])?[0];
579        model.select_output_outlets(&[iff])?;
580
581        let model = model.into_decluttered()?;
582
583        let iff_count = model.nodes().iter().filter(|n| n.op_as::<Iff>().is_some()).count();
584        assert_eq!(iff_count, 0, "Expected Iff to be folded");
585        assert_eq!(
586            model.output_fact(0)?.shape.to_tvec(),
587            tvec![1.to_dim(), 2.to_dim(), 3.to_dim()]
588        );
589        Ok(())
590    }
591}
tract_core/ops/logic.rs

tract_core/ops/
logic.rs
