Skip to content

Commit

Permalink
fix: relax shape checking to >= 0 for onnx
Browse files Browse the repository at this point in the history
  • Loading branch information
ganler committed Oct 10, 2022
1 parent 09c4447 commit 11a29ef
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions nnsmith/materialize/onnx/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,15 +116,15 @@ def analyze_onnx_io(
shape = [dim.dim_value for dim in input.type.tensor_type.shape.dim]
# assert all shapes are positive integers
assert all(
isinstance(dim, int) and dim > 0 for dim in shape
isinstance(dim, int) and dim >= 0 for dim in shape
), f"Fixed shape needed, but got {shape} for input {name}"
input_types[name] = AbsTensor(shape=shape, dtype=dtype)
for output in model.graph.output:
name = output.name
dtype = dtype_from_onnx(output.type.tensor_type.elem_type)
shape = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
assert all(
isinstance(dim, int) and dim > 0 for dim in shape
isinstance(dim, int) and dim >= 0 for dim in shape
), f"Fixed shape needed, but got {shape} for output {name}"
output_types[name] = AbsTensor(shape=shape, dtype=dtype)
return input_types, output_types
Expand Down

0 comments on commit 11a29ef

Please sign in to comment.