1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
use std::fmt;
use std::ops;
mod tree;
pub use self::tree::TDim;
use crate::TractResult;
pub trait DimLike:
Clone
+ Default
+ PartialEq
+ From<usize>
+ ::num_traits::Zero
+ fmt::Debug
+ fmt::Display
+ std::hash::Hash
+ ops::Add<Self, Output = Self>
+ ops::Add<usize, Output = Self>
+ for<'a> ops::Add<&'a Self, Output = Self>
+ ops::Sub<Self, Output = Self>
+ ops::Sub<usize, Output = Self>
+ for<'a> ops::Sub<&'a Self, Output = Self>
+ ops::Mul<usize, Output = Self>
+ ops::Div<usize, Output = Self>
+ ops::Rem<usize, Output = Self>
+ Send
+ Sync
+ 'static
+ std::iter::Sum
+ ToDim
{
fn maybe_mul(&self, other: &Self) -> TractResult<Self>;
fn div_ceil(&self, other: usize) -> Self {
(self.clone() + other - 1) / other
}
fn to_integer(&self) -> TractResult<i32>;
fn one() -> Self;
}
impl DimLike for TDim {
fn maybe_mul(&self, other: &Self) -> TractResult<Self> {
if let Ok(d) = other.to_integer() {
Ok(self.clone() * d)
} else if let Ok(a) = self.to_integer() {
Ok(other.clone() * a)
} else {
bail!("product with too many symbols")
}
}
fn to_integer(&self) -> TractResult<i32> {
TDim::to_integer(self)
}
fn one() -> Self {
Self::from(1)
}
}
impl DimLike for usize {
fn maybe_mul(&self, other: &Self) -> TractResult<Self> {
Ok(self * other)
}
fn to_integer(&self) -> TractResult<i32> {
Ok(*self as i32)
}
fn one() -> usize {
1
}
}
pub trait MaybeProduct<D> {
fn maybe_product(self) -> TractResult<D>;
}
impl<D: DimLike, A: std::borrow::Borrow<D>, I: Iterator<Item = A>> MaybeProduct<D> for I {
fn maybe_product(mut self) -> TractResult<D> {
self.try_fold(D::one(), |acc, d| acc.maybe_mul(d.borrow()))
}
}
pub trait ToDim {
fn to_dim(self) -> TDim;
}
impl<I: Into<TDim>> ToDim for I {
fn to_dim(self) -> TDim {
self.into()
}
}