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

tract_core/ops/
logic.rs
