Skip to content

Commit

Permalink
[torchlib] Improve aten::fill (#1754)
Browse files Browse the repository at this point in the history
I updated torch-onnx to handle empty `[]` inputs, so the isinstance
check is not needed.
  • Loading branch information
justinchuby authored Jul 25, 2024
1 parent 712aa87 commit 937558f
Show file tree
Hide file tree
Showing 3 changed files with 2 additions and 5 deletions.
4 changes: 0 additions & 4 deletions onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -3622,10 +3622,6 @@ def aten_full(

if dtype != -1:
fill_value = op.Cast(fill_value, to=dtype)
if isinstance(size, list) and size == []:
# TODO(justinchuby): Handle empty list better than using isinstance
# size can be empty, meaning a scalar
return fill_value

size = op.Cast(size, to=INT64.dtype)
return op.Expand(fill_value, size)
Expand Down
2 changes: 1 addition & 1 deletion onnxscript/tools/benchmark/benchmark_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ def common_export(
if exporter == "script":
torch.onnx.export(
model,
inputs,
inputs, # type: ignore[arg-type]
filename,
do_constant_folding=False,
input_names=[f"input{i}" for i in range(len(inputs))],
Expand Down
1 change: 1 addition & 0 deletions onnxscript/tools/transformers_models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ def export_to_onnx(
prog = torch.onnx.export(model, args, dynamo=True) # pylint: disable=no-value-for-parameter
else:
prog = torch.onnx.dynamo_export(model, *args)
assert prog is not None
model_proto = prog.model_proto
if optimize:
model_proto = onnxscript.optimizer.optimize(
Expand Down

0 comments on commit 937558f

Please sign in to comment.