oxfuse_core/
shape.rs

1//! Shape rules: one function per concern, all arithmetic checked.
2
3use crate::{Error, Limits, Op, TensorType};
4
5/// Product of the dimensions, or `None` on overflow. Rank 0 gives 1.
6pub fn checked_num_elements(shape: &[usize]) -> Option<usize> {
7    shape.iter().try_fold(1usize, |acc, &d| acc.checked_mul(d))
8}
9
10/// Validates a declared shape against the limits and returns its element count.
11pub fn check_shape(name: &str, shape: &[usize], limits: &Limits) -> Result<usize, Error> {
12    if shape.len() > limits.max_rank {
13        return Err(Error::TensorLimit {
14            name: name.to_owned(),
15            what: "tensor rank",
16            found: shape.len(),
17            limit: limits.max_rank,
18        });
19    }
20    let invalid = |reason: &'static str| Error::InvalidShape {
21        name: name.to_owned(),
22        shape: shape.to_vec(),
23        reason,
24    };
25    if shape.contains(&0) {
26        return Err(invalid("dimensions must be at least 1"));
27    }
28    let count = checked_num_elements(shape).ok_or_else(|| invalid("element count overflows"))?;
29    if count > limits.max_tensor_elements {
30        return Err(Error::TensorLimit {
31            name: name.to_owned(),
32            what: "tensor elements",
33            found: count,
34            limit: limits.max_tensor_elements,
35        });
36    }
37    Ok(count)
38}
39
40fn mismatch(node: &str, reason: String) -> Error {
41    Error::ShapeMismatch {
42        node: node.to_owned(),
43        reason,
44    }
45}
46
47/// The general NumPy broadcasting rule: align trailing dimensions; a
48/// dimension of 1 or a missing dimension stretches to the other operand's.
49pub fn numpy_broadcast(a: &[usize], b: &[usize]) -> Option<Vec<usize>> {
50    let rank = a.len().max(b.len());
51    let dim = |s: &[usize], i: usize| {
52        if i + s.len() >= rank {
53            s[i + s.len() - rank]
54        } else {
55            1
56        }
57    };
58    (0..rank)
59        .map(|i| match (dim(a, i), dim(b, i)) {
60            (x, y) if x == y => Some(x),
61            (1, y) => Some(y),
62            (x, 1) => Some(x),
63            _ => None,
64        })
65        .collect()
66}
67
68/// The subset of `numpy_broadcast` that v1 accepts: equal shapes, or one
69/// operand that is a scalar. Widening this is a change to this one function.
70pub fn elementwise_shape(
71    node: &str,
72    op: &str,
73    a: &[usize],
74    b: &[usize],
75) -> Result<Vec<usize>, Error> {
76    if a == b || b.is_empty() {
77        return Ok(a.to_vec());
78    }
79    if a.is_empty() {
80        return Ok(b.to_vec());
81    }
82    let hint = if numpy_broadcast(a, b).is_some() {
83        " (NumPy broadcasting is not supported yet)"
84    } else {
85        ""
86    };
87    Err(mismatch(
88        node,
89        format!(
90            "{op}: operand shapes {a:?} and {b:?} must be equal or one operand must be a scalar{hint}"
91        ),
92    ))
93}
94
95/// `[M, K] x [K, N] -> [M, N]`. The rule generalizes to `[..., M, K] x
96/// [..., K, N]` with matching batch dimensions; v1 accepts rank 2 only.
97fn matmul_shape(node: &str, a: &[usize], b: &[usize]) -> Result<Vec<usize>, Error> {
98    if a.len() != 2 || b.len() != 2 {
99        return Err(mismatch(
100            node,
101            format!(
102                "MatMul: operands must be rank 2, got ranks {} and {}",
103                a.len(),
104                b.len()
105            ),
106        ));
107    }
108    if a[1] != b[0] {
109        return Err(mismatch(
110            node,
111            format!("MatMul: inner dimensions differ: {a:?} x {b:?}"),
112        ));
113    }
114    Ok(vec![a[0], b[1]])
115}
116
117fn reshape_shape(node: &str, input: &[usize], target: &[usize]) -> Result<Vec<usize>, Error> {
118    match (checked_num_elements(input), checked_num_elements(target)) {
119        (Some(have), Some(want)) if have == want => Ok(target.to_vec()),
120        (Some(have), Some(want)) => Err(mismatch(
121            node,
122            format!(
123                "Reshape: cannot reshape {input:?} ({have} elements) to {target:?} ({want} elements)"
124            ),
125        )),
126        _ => Err(mismatch(
127            node,
128            format!("Reshape: target shape {target:?} element count overflows"),
129        )),
130    }
131}
132
133fn transpose_shape(node: &str, input: &[usize], perm: &[usize]) -> Result<Vec<usize>, Error> {
134    let rank = input.len();
135    if perm.len() != rank {
136        return Err(mismatch(
137            node,
138            format!("Transpose: permutation {perm:?} must have {rank} entries"),
139        ));
140    }
141    let mut seen = vec![false; rank];
142    for &p in perm {
143        if p >= rank || std::mem::replace(&mut seen[p], true) {
144            return Err(mismatch(
145                node,
146                format!("Transpose: {perm:?} is not a permutation of 0..{rank}"),
147            ));
148        }
149    }
150    Ok(perm.iter().map(|&p| input[p]).collect())
151}
152
153/// Infers the output type of `op` applied to `inputs`, checking arity first.
154pub fn infer(node: &str, op: &Op, inputs: &[&TensorType]) -> Result<TensorType, Error> {
155    if inputs.len() != op.arity() {
156        return Err(Error::Arity {
157            node: node.to_owned(),
158            op: op.name(),
159            expected: op.arity(),
160            found: inputs.len(),
161        });
162    }
163    match op {
164        Op::Relu | Op::Sigmoid => Ok(inputs[0].clone()),
165        Op::Transpose { perm } => Ok(TensorType::f32(transpose_shape(
166            node,
167            &inputs[0].shape,
168            perm,
169        )?)),
170        Op::Reshape { shape } => Ok(TensorType::f32(reshape_shape(
171            node,
172            &inputs[0].shape,
173            shape,
174        )?)),
175        Op::MatMul => Ok(TensorType::f32(matmul_shape(
176            node,
177            &inputs[0].shape,
178            &inputs[1].shape,
179        )?)),
180        Op::Add | Op::Mul => Ok(TensorType::f32(elementwise_shape(
181            node,
182            op.name(),
183            &inputs[0].shape,
184            &inputs[1].shape,
185        )?)),
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn rank_zero_has_one_element() {
195        assert_eq!(checked_num_elements(&[]), Some(1));
196    }
197
198    #[test]
199    fn product_overflow_is_none() {
200        assert_eq!(checked_num_elements(&[usize::MAX, 2]), None);
201    }
202
203    #[test]
204    fn check_shape_accepts_valid_shape() {
205        assert_eq!(check_shape("x", &[2, 3], &Limits::default()), Ok(6));
206        assert_eq!(check_shape("s", &[], &Limits::default()), Ok(1));
207    }
208
209    #[test]
210    fn check_shape_rejects_zero_dimension() {
211        let err = check_shape("x", &[2, 0], &Limits::default()).unwrap_err();
212        assert!(matches!(err, Error::InvalidShape { reason, .. } if reason.contains("at least 1")));
213    }
214
215    #[test]
216    fn check_shape_rejects_overflow_without_panicking() {
217        let err = check_shape("x", &[usize::MAX, 2], &Limits::default()).unwrap_err();
218        assert!(matches!(err, Error::InvalidShape { reason, .. } if reason.contains("overflows")));
219    }
220
221    #[test]
222    fn check_shape_rejects_too_many_elements() {
223        let err = check_shape("x", &[20_000, 20_000], &Limits::default()).unwrap_err();
224        assert!(matches!(
225            err,
226            Error::TensorLimit {
227                what: "tensor elements",
228                ..
229            }
230        ));
231    }
232
233    #[test]
234    fn check_shape_rejects_high_rank() {
235        let err = check_shape("x", &[1; 9], &Limits::default()).unwrap_err();
236        assert!(matches!(
237            err,
238            Error::TensorLimit {
239                what: "tensor rank",
240                ..
241            }
242        ));
243    }
244
245    #[test]
246    fn relu_preserves_type() {
247        let t = TensorType::f32(vec![2, 3]);
248        assert_eq!(infer("r", &Op::Relu, &[&t]), Ok(t.clone()));
249    }
250
251    #[test]
252    fn relu_checks_arity() {
253        let err = infer("r", &Op::Relu, &[]).unwrap_err();
254        assert!(matches!(
255            err,
256            Error::Arity {
257                expected: 1,
258                found: 0,
259                ..
260            }
261        ));
262    }
263
264    #[test]
265    fn tensor_limit_errors_name_the_tensor() {
266        let err = check_shape("big", &[20_000, 20_000], &Limits::default()).unwrap_err();
267        assert!(matches!(err, Error::TensorLimit { name, .. } if name == "big"));
268        let err = check_shape("deep", &[1; 9], &Limits::default()).unwrap_err();
269        assert!(matches!(err, Error::TensorLimit { name, .. } if name == "deep"));
270    }
271
272    #[test]
273    fn numpy_broadcast_follows_the_general_rule() {
274        assert_eq!(numpy_broadcast(&[2, 3], &[2, 3]), Some(vec![2, 3]));
275        assert_eq!(numpy_broadcast(&[], &[2, 3]), Some(vec![2, 3]));
276        assert_eq!(numpy_broadcast(&[3], &[2, 3]), Some(vec![2, 3]));
277        assert_eq!(numpy_broadcast(&[2, 1], &[1, 3]), Some(vec![2, 3]));
278        assert_eq!(numpy_broadcast(&[2, 3], &[3, 2]), None);
279        assert_eq!(numpy_broadcast(&[], &[]), Some(vec![]));
280    }
281
282    #[test]
283    fn elementwise_accepts_equal_shapes_and_scalars_only() {
284        assert_eq!(
285            elementwise_shape("n", "Add", &[2, 3], &[2, 3]),
286            Ok(vec![2, 3])
287        );
288        assert_eq!(elementwise_shape("n", "Add", &[], &[2, 3]), Ok(vec![2, 3]));
289        assert_eq!(elementwise_shape("n", "Add", &[2, 3], &[]), Ok(vec![2, 3]));
290        assert_eq!(elementwise_shape("n", "Mul", &[], &[]), Ok(vec![]));
291    }
292
293    #[test]
294    fn elementwise_rejects_incompatible_shapes() {
295        let err = elementwise_shape("y", "Add", &[2, 3], &[3, 2]).unwrap_err();
296        assert_eq!(
297            err.to_string(),
298            "node `y`: Add: operand shapes [2, 3] and [3, 2] must be equal or one operand must be a scalar"
299        );
300    }
301
302    #[test]
303    fn elementwise_names_the_missing_broadcast_feature() {
304        let err = elementwise_shape("y", "Mul", &[2, 3], &[3]).unwrap_err();
305        assert_eq!(
306            err.to_string(),
307            "node `y`: Mul: operand shapes [2, 3] and [3] must be equal or one operand must be a scalar (NumPy broadcasting is not supported yet)"
308        );
309    }
310
311    #[test]
312    fn binary_ops_check_arity() {
313        let t = TensorType::f32(vec![2]);
314        let err = infer("y", &Op::Add, &[&t]).unwrap_err();
315        assert!(matches!(
316            err,
317            Error::Arity {
318                expected: 2,
319                found: 1,
320                ..
321            }
322        ));
323    }
324
325    #[test]
326    fn sigmoid_preserves_type() {
327        let t = TensorType::f32(vec![4]);
328        assert_eq!(infer("s", &Op::Sigmoid, &[&t]), Ok(t.clone()));
329    }
330
331    #[test]
332    fn matmul_shape_is_m_by_n() {
333        let a = TensorType::f32(vec![2, 3]);
334        let b = TensorType::f32(vec![3, 4]);
335        assert_eq!(
336            infer("y", &Op::MatMul, &[&a, &b]),
337            Ok(TensorType::f32(vec![2, 4]))
338        );
339    }
340
341    #[test]
342    fn matmul_rejects_inner_dimension_mismatch() {
343        let a = TensorType::f32(vec![2, 3]);
344        let b = TensorType::f32(vec![4, 5]);
345        let err = infer("y", &Op::MatMul, &[&a, &b]).unwrap_err();
346        assert_eq!(
347            err.to_string(),
348            "node `y`: MatMul: inner dimensions differ: [2, 3] x [4, 5]"
349        );
350    }
351
352    #[test]
353    fn matmul_rejects_non_rank_two_operands() {
354        let a = TensorType::f32(vec![2, 3, 4]);
355        let b = TensorType::f32(vec![4, 5]);
356        let err = infer("y", &Op::MatMul, &[&a, &b]).unwrap_err();
357        assert_eq!(
358            err.to_string(),
359            "node `y`: MatMul: operands must be rank 2, got ranks 3 and 2"
360        );
361    }
362
363    fn reshape(to: Vec<usize>) -> Op {
364        Op::Reshape { shape: to }
365    }
366
367    #[test]
368    fn reshape_keeps_element_count() {
369        let t = TensorType::f32(vec![2, 3]);
370        assert_eq!(
371            infer("r", &reshape(vec![3, 2]), &[&t]),
372            Ok(TensorType::f32(vec![3, 2]))
373        );
374        assert_eq!(
375            infer("r", &reshape(vec![6]), &[&t]),
376            Ok(TensorType::f32(vec![6]))
377        );
378    }
379
380    #[test]
381    fn reshape_handles_scalars_and_identity() {
382        let scalar = TensorType::f32(vec![]);
383        assert_eq!(
384            infer("r", &reshape(vec![1]), &[&scalar]),
385            Ok(TensorType::f32(vec![1]))
386        );
387        let one = TensorType::f32(vec![1]);
388        assert_eq!(
389            infer("r", &reshape(vec![]), &[&one]),
390            Ok(TensorType::f32(vec![]))
391        );
392        let t = TensorType::f32(vec![2, 3]);
393        assert_eq!(infer("r", &reshape(vec![2, 3]), &[&t]), Ok(t.clone()));
394    }
395
396    #[test]
397    fn reshape_rejects_element_count_change() {
398        let t = TensorType::f32(vec![2, 3]);
399        let err = infer("y", &reshape(vec![4]), &[&t]).unwrap_err();
400        assert_eq!(
401            err.to_string(),
402            "node `y`: Reshape: cannot reshape [2, 3] (6 elements) to [4] (4 elements)"
403        );
404    }
405
406    #[test]
407    fn reshape_target_overflow_is_rejected() {
408        let t = TensorType::f32(vec![2, 3]);
409        let err = infer("y", &reshape(vec![usize::MAX, 2]), &[&t]).unwrap_err();
410        assert!(matches!(err, Error::ShapeMismatch { reason, .. } if reason.contains("overflows")));
411    }
412
413    fn transpose(perm: Vec<usize>) -> Op {
414        Op::Transpose { perm }
415    }
416
417    #[test]
418    fn transpose_permutes_dimensions() {
419        let t = TensorType::f32(vec![2, 3, 4]);
420        assert_eq!(
421            infer("t", &transpose(vec![2, 0, 1]), &[&t]),
422            Ok(TensorType::f32(vec![4, 2, 3]))
423        );
424    }
425
426    #[test]
427    fn transpose_handles_low_ranks_and_size_one_dims() {
428        let scalar = TensorType::f32(vec![]);
429        assert_eq!(
430            infer("t", &transpose(vec![]), &[&scalar]),
431            Ok(scalar.clone())
432        );
433        let vector = TensorType::f32(vec![5]);
434        assert_eq!(
435            infer("t", &transpose(vec![0]), &[&vector]),
436            Ok(vector.clone())
437        );
438        let ones = TensorType::f32(vec![1, 3, 1]);
439        assert_eq!(
440            infer("t", &transpose(vec![1, 2, 0]), &[&ones]),
441            Ok(TensorType::f32(vec![3, 1, 1]))
442        );
443    }
444
445    #[test]
446    fn invalid_permutations_are_rejected() {
447        let t = TensorType::f32(vec![2, 3]);
448        let duplicate = infer("y", &transpose(vec![0, 0]), &[&t]).unwrap_err();
449        assert_eq!(
450            duplicate.to_string(),
451            "node `y`: Transpose: [0, 0] is not a permutation of 0..2"
452        );
453        let out_of_range = infer("y", &transpose(vec![0, 2]), &[&t]).unwrap_err();
454        assert_eq!(
455            out_of_range.to_string(),
456            "node `y`: Transpose: [0, 2] is not a permutation of 0..2"
457        );
458        let wrong_length = infer("y", &transpose(vec![0]), &[&t]).unwrap_err();
459        assert_eq!(
460            wrong_length.to_string(),
461            "node `y`: Transpose: permutation [0] must have 2 entries"
462        );
463    }
464
465    #[test]
466    fn check_shape_accepts_exactly_the_rank_and_element_limits() {
467        let limits = Limits::default();
468        assert_eq!(check_shape("x", &[1; 8], &limits), Ok(1));
469        assert_eq!(
470            check_shape("x", &[limits.max_tensor_elements], &limits),
471            Ok(limits.max_tensor_elements)
472        );
473        assert!(check_shape("x", &[limits.max_tensor_elements + 1], &limits).is_err());
474    }
475}
oxfuse_core/shape.rs

oxfuse_core/
shape.rs
