tract_core/ops/logic/
comparison.rs

1use crate::broadcast::multi_broadcast;
2use crate::internal::*;
3use crate::ndarray::Zip;
4use crate::ops::binary::BinMiniOp;
5
6use tract_data::TooEarly;
7
8// Helper for eval_out_of_place dispatch
9fn eval_comp_oop<T: Datum + PartialOrd>(
10    a: &Tensor,
11    b: &Tensor,
12    f: impl Fn(&T, &T) -> bool,
13) -> TractResult<Tensor> {
14    let a = a.to_plain_array_view::<T>()?;
15    let b = b.to_plain_array_view::<T>()?;
16    let shape = multi_broadcast(&[a.shape(), b.shape()])?;
17    let mut c = unsafe { Tensor::uninitialized::<bool>(&shape)? };
18    let mut c_plain = c.try_as_plain_ram_mut()?;
19    let mut view = c_plain.to_array_view_mut::<bool>()?;
20    Zip::from(&mut view).and_broadcast(&a).and_broadcast(&b).for_each(|c, a, b| *c = f(a, b));
21    Ok(c)
22}
23
24// Helper for TDim symbolic eval
25fn eval_tdim_symbolic(
26    ctx: &EvalContext,
27    inputs: &TVec<TValue>,
28    prove: impl Fn(&TDim, &TDim) -> TractResult<bool>,
29) -> TractResult<Option<TVec<TValue>>> {
30    rule_if!(
31        inputs[0].datum_type() == TDim::datum_type()
32            || inputs[1].datum_type() == TDim::datum_type()
33    );
34    // a dynamic Range yields i64 while a TDim konst stays TDim: cast the
35    // integer side up so the pair goes through the symbolic path
36    let to_tdim = |t: &TValue| -> TractResult<Tensor> {
37        Ok(t.clone().into_tensor().cast_to_dt(TDim::datum_type())?.into_owned())
38    };
39    let mut a = to_tdim(&inputs[0])?;
40    let mut b = to_tdim(&inputs[1])?;
41    for a in a.try_as_plain_ram_mut()?.as_slice_mut::<TDim>()? {
42        *a = a.eval(ctx.symbols);
43    }
44    for b in b.try_as_plain_ram_mut()?.as_slice_mut::<TDim>()? {
45        *b = b.eval(ctx.symbols);
46    }
47    if let (Ok(a_i64), Ok(b_i64)) = (a.cast_to::<i64>(), b.cast_to::<i64>()) {
48        let result = eval_comp_oop::<i64>(&a_i64, &b_i64, |a, b| {
49            prove(&(*a).into(), &(*b).into()).unwrap_or(false)
50        })?;
51        return Ok(Some(tvec!(result.into_tvalue())));
52    }
53    let a_view = a.to_plain_array_view::<TDim>()?;
54    let b_view = b.to_plain_array_view::<TDim>()?;
55    let shape = multi_broadcast(&[a_view.shape(), b_view.shape()])?;
56    let mut c = unsafe { Tensor::uninitialized::<bool>(&shape)? };
57    let mut c_plain = c.try_as_plain_ram_mut()?;
58    let mut view = c_plain.to_array_view_mut::<bool>()?;
59    let a_bc = a_view.broadcast(&*shape).unwrap();
60    let b_bc = b_view.broadcast(&*shape).unwrap();
61    for ixs in tract_ndarray::indices(&*shape) {
62        view[&ixs] = prove(&a_bc[&ixs], &b_bc[&ixs])?;
63    }
64    Ok(Some(tvec!(c.into_tvalue())))
65}
66
67macro_rules! comp_bin_mini_op {
68    ($Op:ident, $name:literal, $cmp:tt, $prove_tdim:expr, $uniform_tdim:expr) => {
69        #[derive(Debug, Clone, Hash, PartialEq, Eq)]
70        pub struct $Op;
71
72        impl BinMiniOp for $Op {
73            fn name(&self) -> &'static str {
74                $name
75            }
76
77            fn result_datum_type(&self, _a: DatumType, _b: DatumType) -> TractResult<DatumType> {
78                Ok(bool::datum_type())
79            }
80
81            fn is_commutative(&self) -> bool {
82                false
83            }
84
85            fn eval_in_a(&self, _a: &mut Tensor, _b: &Tensor) -> TractResult<()> {
86                bail!("Comparison changes datum type, eval_in_a not supported")
87            }
88
89            fn eval_out_of_place(
90                &self,
91                c: &mut Tensor,
92                a: &Tensor,
93                b: &Tensor,
94            ) -> TractResult<()> {
95                let dt = a.datum_type();
96                if dt == String::datum_type() {
97                    let a = a.to_plain_array_view::<String>()?;
98                    let b = b.to_plain_array_view::<String>()?;
99                    let mut c_plain = c.try_as_plain_ram_mut()?;
100                    let mut view = c_plain.to_array_view_mut::<bool>()?;
101                    Zip::from(&mut view).and_broadcast(&a).and_broadcast(&b)
102                        .for_each(|c, a, b| *c = a $cmp b);
103                    return Ok(());
104                }
105                fn inner<T: Datum + PartialOrd>(c: &mut Tensor, a: &Tensor, b: &Tensor, f: impl Fn(&T, &T) -> bool) -> TractResult<()> {
106                    let a = a.to_plain_array_view::<T>()?;
107                    let b = b.to_plain_array_view::<T>()?;
108                    let mut c_plain = c.try_as_plain_ram_mut()?;
109                    let mut view = c_plain.to_array_view_mut::<bool>()?;
110                    Zip::from(&mut view).and_broadcast(&a).and_broadcast(&b)
111                        .for_each(|c, a, b| *c = f(a, b));
112                    Ok(())
113                }
114                dispatch_numbers!(inner(dt)(c, a, b, |a: &_, b: &_| a $cmp b))
115            }
116
117            fn eval(&self, a: TValue, b: TValue, c_dt: DatumType) -> TractResult<Tensor> {
118                let c_shape = crate::broadcast::multi_broadcast(&[a.shape(), b.shape()])?;
119                let mut c = unsafe { Tensor::uninitialized_dt(c_dt, &c_shape)? };
120                self.eval_out_of_place(&mut c, &a, &b)?;
121                Ok(c)
122            }
123
124            fn eval_symbolic(
125                &self,
126                ctx: &EvalContext,
127                inputs: TVec<TValue>,
128            ) -> TractResult<Option<TVec<TValue>>> {
129                eval_tdim_symbolic(ctx, &inputs, $prove_tdim)
130            }
131
132            fn uniform_tdim_comparison(
133                &self,
134                a: &TDim,
135                b: &TDim,
136            ) -> Option<TDim> {
137                Some(($uniform_tdim)(a, b))
138            }
139        }
140    };
141}
142
143fn prove_eq(a: &TDim, b: &TDim) -> TractResult<bool> {
144    Ok(a == b)
145}
146
147fn prove_ne(a: &TDim, b: &TDim) -> TractResult<bool> {
148    Ok(a != b)
149}
150
151fn prove_gte(a: &TDim, b: &TDim) -> TractResult<bool> {
152    let diff = a.clone() - b;
153    if diff.prove_positive_or_zero() {
154        Ok(true)
155    } else if diff.prove_strict_negative() {
156        Ok(false)
157    } else {
158        bail!(TooEarly::UndeterminedSymbol(diff.to_string()))
159    }
160}
161
162fn prove_gt(a: &TDim, b: &TDim) -> TractResult<bool> {
163    let diff = a.clone() - b;
164    if diff.prove_strict_positive() {
165        Ok(true)
166    } else if diff.prove_negative_or_zero() {
167        Ok(false)
168    } else {
169        bail!(TooEarly::UndeterminedSymbol(diff.to_string()))
170    }
171}
172
173fn prove_lte(a: &TDim, b: &TDim) -> TractResult<bool> {
174    prove_gte(b, a)
175}
176
177fn prove_lt(a: &TDim, b: &TDim) -> TractResult<bool> {
178    prove_gt(b, a)
179}
180
181comp_bin_mini_op!(CompEq, "Eq", ==, prove_eq, |a: &TDim, b: &TDim|
182    TDim::Eq(Box::new(a.clone()), Box::new(b.clone())).reduce()
183);
184
185comp_bin_mini_op!(CompNE, "NE", !=, prove_ne, |a: &TDim, b: &TDim|
186    (TDim::Val(1) - TDim::Eq(Box::new(a.clone()), Box::new(b.clone()))).reduce()
187);
188
189comp_bin_mini_op!(CompLT, "LT", <, prove_lt, |a: &TDim, b: &TDim|
190    TDim::Ge(Box::new(b.clone()), Box::new((a.clone() + TDim::Val(1)).reduce())).reduce()
191);
192
193comp_bin_mini_op!(CompGT, "GT", >, prove_gt, |a: &TDim, b: &TDim|
194    TDim::Ge(Box::new((a.clone() + TDim::Val(1)).reduce()), Box::new(b.clone())).reduce()
195);
196
197comp_bin_mini_op!(CompLTE, "LTE", <=, prove_lte, |a: &TDim, b: &TDim|
198    TDim::Ge(Box::new(b.clone()), Box::new(a.clone())).reduce()
199);
200
201comp_bin_mini_op!(CompGTE, "GTE", >=, prove_gte, |a: &TDim, b: &TDim|
202    TDim::Ge(Box::new(a.clone()), Box::new(b.clone())).reduce()
203);
204
205// Factory functions
206pub fn comp_eq() -> Box<dyn BinMiniOp> {
207    Box::new(CompEq)
208}
209pub fn comp_ne() -> Box<dyn BinMiniOp> {
210    Box::new(CompNE)
211}
212pub fn comp_lt() -> Box<dyn BinMiniOp> {
213    Box::new(CompLT)
214}
215pub fn comp_gt() -> Box<dyn BinMiniOp> {
216    Box::new(CompGT)
217}
218pub fn comp_lte() -> Box<dyn BinMiniOp> {
219    Box::new(CompLTE)
220}
221pub fn comp_gte() -> Box<dyn BinMiniOp> {
222    Box::new(CompGTE)
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    // A dynamic Range output (i64) compared against a TDim konst must evaluate
230    // through the symbolic path (the i64 side is cast up), not error on the
231    // mixed datum types.
232    #[test]
233    fn mixed_i64_tdim_comparison() -> TractResult<()> {
234        use crate::ops::array::Range;
235        use crate::ops::binary::TypedBinOp;
236
237        let mut model = TypedModel::default();
238        let s = model.symbols.sym("S");
239        let definer = model.add_source("definer", f32::datum_type().fact([s.to_dim()]))?;
240        let start = model.add_const("start", tensor0(TDim::Val(0)))?;
241        let step = model.add_const("step", tensor0(TDim::Val(1)))?;
242        let end = model.wire_node(
243            "end",
244            crate::ops::konst::Const::new(tensor0(TDim::Sym(s.clone())).into_arc_tensor())?,
245            &[],
246        )?;
247        let range = model.wire_node("range", Range::new(s.to_dim()), &[start, end[0], step])?;
248        let t_const = model.wire_node(
249            "t_const",
250            crate::ops::konst::Const::new(tensor0(TDim::Sym(s.clone())).into_arc_tensor())?,
251            &[],
252        )?;
253        let t_unsq =
254            model.wire_node("t_unsq", crate::ops::change_axes::AxisOp::Add(0), &[t_const[0]])?;
255        let lt = model.wire_node("lt", TypedBinOp(comp_lt(), None), &[range[0], t_unsq[0]])?;
256        model.select_output_outlets(&[lt[0], definer])?;
257        let found = crate::internal::TypedSimplePlan::new(model)?
258            .run(tvec!(tensor1(&[0f32; 3]).into_tvalue()))?;
259        // all of range(0..S) is strictly below S: true, true, true
260        assert_eq!(*found[0], tensor1(&[true, true, true]));
261        Ok(())
262    }
263}
tract_core/ops/logic/comparison.rs

tract_core/ops/logic/
comparison.rs
