Skip to content

Commit

Permalink
Correctly mark panic conditions in op_configuration
Browse files Browse the repository at this point in the history
  • Loading branch information
Dirleye committed Jul 2, 2024
1 parent b169bfe commit e755942
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions crates/burn-import/src/onnx/op_configuration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1002,22 +1002,22 @@ pub fn reduce_prod_config(node: &Node) -> Option<usize> {
match key.as_str() {
"axes" => axes = value.clone().into_i64s(),
"keepdims" => keepdims = value.clone().into_i64(),
// TODO: handle noop_with_empty_axes
// TODO: handle noop_with_empty_axes (opset 18)
_ => {}
}
}

if axes.len() > 1 {
panic!("ReduceMean: reducing on multiple dimensions is not supported")
panic!("ReduceProd: reducing on multiple dimensions is not supported")
}

if axes.is_empty() && keepdims == 1 {
panic!("ReduceMean: axes must be provided with keepdims")
panic!("ReduceProd: axes must be provided with keepdims")
}

if !axes.is_empty() && keepdims == 0 {
// Not supported in Burn
panic!("ReduceMean: the reduce operation must preserve the reduced dimension")
panic!("ReduceProd: the reduce operation must preserve the reduced dimension")
}

if axes.is_empty() {
Expand Down

0 comments on commit e755942

Please sign in to comment.