Repository: pytorch/pytorch

Issue:
[ONNX] Fix aten::new_zeros export
Fixes export in https://github.com/pytorch/pytorch/issues/110597

cc @justinchuby 

[PyTorch 2.1 regression] TorchScript behavior changed from 2.0.1 (and older) to 2.1
### 🐛 Describe the bug

The ONNX export caught a behavior change on TorchScript API after PyTorch 2.1 was released. Before, a `Torch._C.Value.node.mustBeNone()` returned `False` for a model with `aten::new_zeros` ops while after PyTorch 2.1 it returns `True`, changing the execution path of the ONNX model export.

There is nothing in the [Pytorch 2.1 Release Notes](https://github.com/pytorch/pytorch/releases/tag/v2.1.0) that mentions a behavior change of this nature

Reproduction:
```
import torch.nn as nn
import torch
import onnxruntime as ort

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(50, 64)
    
    def forward(self, x):
        inp = x.new_zeros(x.shape)
        
        return self.emb(inp)

model = MyModel()
inp = torch.Tensor([[2, 5, 6], [3, 2, 5]]).to(torch.int64)

torch.onnx.export(model, (inp,), "model.onnx", opset_version=9)
```

The repro uses `torch.onnx.export` because it calls `torch.jit.trace` and provides an easy entry point to print `new_zeros` arguments. 

The ONNX error is as shown:

```bash
/opt/pytorch/torch/onnx/utils.py:1703: UserWarning: The exported ONNX model failed ONNX shape inference. The model will not be executable by the ONNX Runtime. If this is unintended and you believe there is a bug, please report an issue at https://github.com/pytorch/pytorch/issues. Error reported by strict ONNX shape inference: [ShapeInferenceError] (op_type:Gather, node name: /emb/Gather): indices typestr: Tind, has unsupported type: tensor(float) (Triggered internally at /opt/pytorch/torch/csrc/jit/serialization/export.cpp:1445.)
  _C._check_onnx_proto(proto)
Traceback (most recent call last):
  File "repro_jit_type.py", line 18, in <module>
    torch.onnx.export(model, (inp,), "model.onnx", opset_version=9)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 419, in __init__
    self._create_inference_session(providers, provider_options, disabled_optimizers)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 460, in _create_inference_session
    sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from model.onnx failed:This is an invalid model. Type Error: Type 'tensor(float)' of input parameter (/Constant_output_0) of operator (Gather) in node (/emb/Gather) is invalid.
```

The error above doesn't point to TorchScript, but shows the execution path has changed for ONNX exporter.
Looking into it further and printing the `new_zeros`'s `dtype` argument (from within the [ONNX exporter code](https://github.com/pytorch/pytorch/blob/144cda7f068854dd870c9567781aa2aca6d5e4cf/torch/onnx/symbolic_opset9.py#L388) which has an easy entry point to the `new_zeros` TorchScript node), we can see that on PyTorch 2.1, the `dtype` arg of `new_zeros` is captured as:

```python
> /opt/pytorch/torch/onnx/symbolic_opset9.py(3889)new_zeros()
-> self_dtype = symbolic_helper._try_get_scalar_type(self)
(Pdb) dtype
15 defined in (%15 : NoneType = prim::Constant(), scope: __main__.MyModel::
)
```

However, in PyTorch 2.0.1 and older, it was:

```python
> /opt/conda/envs/ptca/lib/python3.8/site-packages/torch/onnx/symbolic_opset9.py(3708)new_zeros()
-> self_dtype = symbolic_helper._try_get_scalar_type(self)
(Pdb) dtype
15 defined in (%15 : Long(device=cpu) = onnx::Constant[value={4}](), scope: __main__.MyModel::
)
```

### Versions

PyTorch 2.1

cc @ezyang @gchanan @zou3519 @kadeng @EikanWang @jgong5 @wenzhe-nrv @sanchitintel


Discussion:
<!-- drci-comment-start -->

## :link: Helpful Links
### :test_tube: See artifacts and rendered test results at [hud.pytorch.org/pr/110935](https://hud.pytorch.org/pr/110935)
* :page_facing_up: Preview [Python docs built from this PR](https://docs-preview.pytorch.org/pytorch/pytorch/110935/index.html)
* :page_facing_up: Preview [C++ docs built from this PR](https://docs-preview.pytorch.org/pytorch/pytorch/110935/cppdocs/index.html)
* :question: Need help or want to give feedback on the CI? Visit the [bot commands wiki](https://github.com/pytorch/pytorch/wiki/Bot-commands) or our [office hours](https://github.com/pytorch/pytorch/wiki/Dev-Infra-Office-Hours)

Note: Links to docs will display an error until the docs builds have been completed.


## :x: 1 New Failure
As of commit 5d318d5fcca6e06317a0ce9c248b89ebb194caa5 with merge base 37567fdf311549dedf74731b359edc6b077f30a0 (<sub><sub><img alt="image" width=70 src="https://img.shields.io/date/1696925223?label=&color=FFFFFF&style=flat-square"></sub></sub>):
<details open><summary><b>NEW FAILURE</b> - The following job has failed:</summary><p>

* [pull / linux-focal-py3.8-clang10-onnx / test (default, 1, 2, linux.2xlarge)](https://hud.pytorch.org/pr/pytorch/pytorch/110935#17568957963) ([gh](https://github.com/pytorch/pytorch/actions/runs/6470973050/job/17568957963))
</p></details>


This comment was automatically generated by Dr. CI and updates every 15 minutes.
<!-- drci-comment-end -->
<a href="https://easycla.lfx.linuxfoundation.org/#/?version=2"><img src="https://s3.amazonaws.com/cla-project-logo-prod/cla-signed.svg" alt="CLA Signed" align="left" height="28" width="328" ></a><br/><br />The committers listed above are authorized under a signed CLA.<ul><li>:white_check_mark: login: fxmarty  (4e9117f0c5b9d05e84f06cb80f2b4beeb7e7a7ac, 5d318d5fcca6e06317a0ce9c248b89ebb194caa5)</li></ul>
Could you enable the test in test_op_consistency and fix CLA? Thanks!
Thank you I added the test in test_op_consistency (which is not passing - independent of this PR, happy to address any suggestion here) and signed CLA.
@fxmarty I am testing https://github.com/pytorch/pytorch/pull/110956  which pushes your fix to `from_value` while we don't hear from PyTorch

What is different from the original PR:

* Generic code: The `torch._C.Value.node().mustBeNone()` is encapsulated into the high level API `JitScalarType.from_value` so users don't have to worry about it
* Specific to new_zeros: When checking `dtype`, we always must use `JitScalarType` for all arguments. `new_zeros` calls `from_value` for `self` argument but not for `dtype` too, which skips all the handling provided by `JitScalarType`. Other operators had the same issue and were fixed here too
> Could you enable the test in test_op_consistency and fix CLA? Thanks!

@justinchuby adding `new_zeros` there raised those errors. Any idea on why?

```bash
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool SUBFAIL [0.0000s] [  0%]
onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool PASSED [0.1885s] [  0%]

=================================== FAILURES ===================================
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor(False), ())', kwargs='{}', opset=9, sample_num=0) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array(False)]. kw_inputs: {}. input names: [].
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor([[ True,  True, False,  True,  True],\n        [ True,  True,  True,  True,  True],\n        [False, False,  True, False, False],\n        [False, False, False,  True, False],\n        [ True,  True, False, False,  True]]), (2, 0))', kwargs='{}', opset=9, sample_num=1) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array([[ True,  True, False,  True,  True],
       [ True,  True,  True,  True,  True],
       [False, False,  True, False, False],
       [False, False, False,  True, False],
       [ True,  True, False, False,  True]]), array(2), array(0)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2'].
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor([], size=(0, 5, 0), dtype=torch.bool), (3, 2, 2))', kwargs='{}', opset=9, sample_num=2) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array([], shape=(0, 5, 0), dtype=bool), array(3), array(2), array(2)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2', 'onnx::Unsqueeze_3'].
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor([ True,  True,  True, False,  True]), (2, 3))', kwargs="{'dtype': torch.bool, 'device': 'cpu'}", opset=9, sample_num=3) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array([ True,  True,  True, False,  True]), array(2), array(3)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2'].
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor([False,  True, False,  True,  True]), (10,))', kwargs="{'dtype': torch.float64}", opset=9, sample_num=4) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array([False,  True, False,  True,  True]), array(10)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1'].
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor([False,  True,  True, False, False]), (1, 1, 12))', kwargs="{'device': 'cpu'}", opset=9, sample_num=5) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array([False,  True,  True, False, False]), array(1), array(1), array(12)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2', 'onnx::Unsqueeze_3'].
_ TestOnnxModelOutputConsistency_opset9CPU.test_output_match_new_zeros_cpu_bool (inputs='(tensor([ True, False,  True,  True,  True]), (2, 2, 2))', kwargs="{'dtype': torch.float64, 'device': 'cpu'}", opset=9, sample_num=6) _
Traceback (most recent call last):
  File "onnx/test_op_consistency.py", line 317, in test_output_match
    self.run_test(model, inputs, rtol=rtol, atol=atol)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 204, in run_test
    _run_test(model, tracing_remained_onnx_input_idx)
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 164, in _run_test
    return run_model_test(
  File "/var/lib/jenkins/workspace/test/onnx/onnx_test_common.py", line 84, in run_model_test
    return verification.verify(*args, options=options, **kwargs)
  File "<@beartype(torch.onnx.verification.verify) at 0x7f2ef0a3e160>", line 352, in verify
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 890, in verify
    _compare_onnx_pytorch_model(
  File "<@beartype(torch.onnx.verification._compare_onnx_pytorch_model) at 0x7f2ef0a271f0>", line 160, in _compare_onnx_pytorch_model
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 451, in _compare_onnx_pytorch_model
    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 443, in compare_onnx_pytorch_model_with_input
    onnx_outs = _run_onnx(onnx_session, onnx_inputs)
  File "<@beartype(torch.onnx.verification._run_onnx) at 0x7f2ef0a994c0>", line 16, in _run_onnx
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/onnx/verification.py", line 172, in _run_onnx
    raise ValueError(
ValueError: got too many positional inputs. inputs: [array([ True, False,  True,  True,  True]), array(2), array(2), array(2)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2', 'onnx::Unsqueeze_3'].
- generated xml file: /var/lib/jenkins/workspace/test/test-reports/python-pytest/onnx.test_op_consistency/onnx.test_op_consistency-38452ac968537a1e.xml -
=========================== short test summary info ============================
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array(False)]. kw_inputs: {}. input names: [].
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array([[ True,  True, False,  True,  True],
       [ True,  True,  True,  True,  True],
       [False, False,  True, False, False],
       [False, False, False,  True, False],
       [ True,  True, False, False,  True]]), array(2), array(0)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2'].
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array([], shape=(0, 5, 0), dtype=bool), array(3), array(2), array(2)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2', 'onnx::Unsqueeze_3'].
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array([ True,  True,  True, False,  True]), array(2), array(3)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2'].
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array([False,  True, False,  True,  True]), array(10)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1'].
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array([False,  True,  True, False, False]), array(1), array(1), array(12)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2', 'onnx::Unsqueeze_3'].
/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/__init__.py:635: UserWarning: torch.set_default_tensor_type() is deprecated as of PyTorch 2.1, please use torch.set_default_dtype() and torch.set_default_device() as alternatives. (Triggered internally at /var/lib/jenkins/workspace/torch/csrc/tensor/python_tensor.cpp:453.)
  _C._set_default_tensor_type(t)
SUBFAIL [0.0000s] onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset9CPU::test_output_match_new_zeros_cpu_bool - ValueError: got too many positional inputs. inputs: [array([ True, False,  True,  True,  True]), array(2), array(2), array(2)]. kw_inputs: {}. input names: ['onnx::Unsqueeze_1', 'onnx::Unsqueeze_2', 'onnx::Unsqueeze_3'].
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 7 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
================== 7 failed, 1 passed, 62 deselected in 0.31s ==================
```
I am not very sure. It looks like the exported model was expected 3 inputs but the test provided 4.
Potentially, you may xfail the bool dtypes with an `xfail(..., dtypes=(torch.bool,))` entry to unblock this.
> I am not very sure. It looks like the exported model was expected 3 inputs but the test provided 4.

I think I know why. The tests for `aten::new_zeros` are exported without inputs, because all inputs are probably folded into graph constants. To fix that, the tests should run with `dynamic_axes` specified in the `torch.onnx.export` call
https://github.com/pytorch/pytorch/pull/110956 is being merged to fix this. @fxmarty has been added as co-author on the new PR
[model.onnx.zip](https://github.com/pytorch/pytorch/files/12819009/model.onnx.zip)

Adding the failing model exported from `torch.onnx.export` on 2.1 for reference
Hi @justinchuby @BowenBao @thiagocrepaldi, this issue is related to the dtype `indices` is exported to.

Minimal reproduction:

```python
import torch.nn as nn
import torch
import onnxruntime as ort

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(50, 64)
    
    def forward(self, x):
        inp = x.new_zeros(x.shape)
        
        return self.emb(inp)

model = MyModel()
inp = torch.Tensor([[2, 5, 6], [3, 2, 5]]).to(torch.int64)

torch.onnx.export(model, (inp,), "model.onnx", opset_version=9)

session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
```

Debugging, it appears that the issue is due to this where `dtype` is not None
https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/onnx/symbolic_opset9.py#L3707-L3708

and when calling `zeros` and in turn its decorator `symbolic_helper.parse_args`, we used to go into this controlflow for the argument `dtype`,

https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/onnx/symbolic_helper.py#L82-L83

while in 2.1.0 we go into this controlflow

https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/onnx/symbolic_helper.py#L78-L79

This in the `zeros` function, with 2.0.1 we rightfully call `scalar_type = _type_utils.JitScalarType(dtype)` while in 2.1.0 we call `scalar_type = _type_utils.JitScalarType.FLOAT`.

https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/onnx/symbolic_opset9.py#L3690-L3693

It is unclear to me why `mustBeNone` for this node has changed between 2.0.1 and 2.1, do you have any idea? The node is `15 defined in (%15 : NoneType = prim::Constant(), scope: __main__.MyModel::)` (printed with 2.1.0) with the above example.

Maybe the issue is with torch.jit.trace and not ONNX?

Thank you!
Actually yes, the issue is in the captured `ScriptModule`. Now in 2.1.0 the `dtype` arg of `new_zeros` is captured as `%17 : NoneType = prim::Constant()`, while it used to be capted as `%17 : int = prim::Constant[value=4]() # <tmp 4>+34:11:0`

An easy solution is to replace this check https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/onnx/symbolic_opset9.py#L3707-L3708

by 

```python
    if dtype.node().mustBeNone() and self_dtype is not None:
        dtype = self_dtype
```

Is that fine or is the change in `torch.jit.trace` a bug and the issue should be fixed upstream?
Thanks for discovering the issue and for the detailed analysis! Whether it is a bug in jit can be investigated (from briefly scanning the code it seems fine for dtype to be none). But the change in exporter makes sense to me. 
Thank you. I'll do the change for `new_zeros` then, though the change in jit.trace may impact (or may not) other op as well.
> Actually yes, the issue is in the captured `ScriptModule`. Now in 2.1.0 the `dtype` arg of `new_zeros` is captured as `%17 : NoneType = prim::Constant()`, while it used to be capted as `%17 : int = prim::Constant[value=4]() # <tmp 4>+34:11:0`
> 
> An easy solution is to replace this check
> 
> https://github.com/pytorch/pytorch/blob/e9ebda29d87ce0916ab08c06ab26fd3766a870e5/torch/onnx/symbolic_opset9.py#L3707-L3708
> 
> by
> 
> ```python
>     if dtype.node().mustBeNone() and self_dtype is not None:
>         dtype = self_dtype
> ```
> 
> Is that fine or is the change in `torch.jit.trace` a bug and the issue should be fixed upstream?

@ezyang do you know someone that could help us track a potential regression on TorchScript?
@fxmarty I have updated the issue description to capture the root cause you have identified in a previous [comment](https://github.com/pytorch/pytorch/issues/110597#issuecomment-1750488721)

Feel free to update it accordingly. The important is to show that the ONNX exporter detected the behavior change, not caused it so that TorchScript friends can help us here
cc @davidberard98 can you take a look at the priority of this please?
Tentative high pri for now due to regression
I wonder if difference is caused by dynamic shapes?
Seems unlikely? This is all TS...
The `dtype`, `layout` & `device` arguments of `new_zeros` are being traced as `NoneType`. I'll report back on Friday. Thanks!
was curious where this came from and ran a bisect. Bisect is still running but I'm pretty sure it's from this PR, which basically makes this exact change: https://github.com/pytorch/pytorch/pull/97564.

will take a look at the linked PRs in the next day or two.

@albanD not familiar with priority levels - perhaps you or someone else from core can judge based on the change in the referenced PR?
cc @bdhirsh is your PR expected to impact all forms of tracing this way?
@albanD yep - the purpose of that PR was to avoid burning in the defaults of `*_like` factory functions during tracing. And this burning in previously happened at the python binding layer. So all traces that live below the python layer (make_fx, and also torchscript tracing and friends) will be affected by it.

In theory, the semantics of `aten.ones_like(x)` and `aten.ones_like(x, device='cpu')` are the same. But I should have marked that PR as potentially BC-breaking for backends / consumers of that IR 😕.
> @albanD yep - the purpose of that PR was to avoid burning in the defaults of `*_like` factory functions during tracing. And this burning in previously happened at the python binding layer. So all traces that live below the python layer (make_fx, and also torchscript tracing and friends) will be affected by it.
> 
> In theory, the semantics of `aten.ones_like(x)` and `aten.ones_like(x, device='cpu')` are the same. But I should have marked that PR as potentially BC-breaking for backends / consumers of that IR 😕.

Thank you. 
Are we considering a change to make this a BC again (maybe for 2.1.1) or the new breaking behavior will stand as is?

Resolving patch:
diff --git a/torch/onnx/_type_utils.py b/torch/onnx/_type_utils.py
--- a/torch/onnx/_type_utils.py
+++ b/torch/onnx/_type_utils.py
@@ -170,7 +170,9 @@ def from_value(
             SymbolicValueError: when value.type()'s info are empty and default is None
         """
 
-        if not isinstance(value, (torch._C.Value, torch.Tensor)):
+        if not isinstance(value, (torch._C.Value, torch.Tensor)) or (
+            isinstance(value, torch._C.Value) and value.node().mustBeNone()
+        ):
             # default value of type JitScalarType is returned when value is not valid
             if default is None:
                 raise errors.OnnxExporterError(
diff --git a/torch/onnx/symbolic_helper.py b/torch/onnx/symbolic_helper.py
--- a/torch/onnx/symbolic_helper.py
+++ b/torch/onnx/symbolic_helper.py
@@ -476,8 +476,8 @@ def _if_scalar_type_as(self, tensor):
 
 
 @_beartype.beartype
-def _is_none(x: _C.Value) -> bool:
-    return x.node().mustBeNone()
+def _is_none(x: Any) -> bool:
+    return x is None or (x.node().mustBeNone() if isinstance(x, _C.Value) else False)
 
 
 @_beartype.beartype
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
--- a/torch/onnx/symbolic_opset9.py
+++ b/torch/onnx/symbolic_opset9.py
@@ -3785,7 +3785,7 @@ def new_empty(
     g: jit_utils.GraphContext, self, sizes, dtype, layout, device, pin_memory=False
 ):
     self_dtype = symbolic_helper._try_get_scalar_type(self)
-    if dtype is None and self_dtype is not None:
+    if symbolic_helper._is_none(dtype) and self_dtype is not None:
         dtype = self_dtype
     return empty(g, sizes, dtype, layout, device, pin_memory)
 
@@ -3867,7 +3867,7 @@ def zeros_like(
     memory_format=None,
 ):
     shape = g.op("Shape", input)
-    if dtype is None:
+    if symbolic_helper._is_none(dtype):
         scalar_type = _type_utils.JitScalarType.from_value(
             input, _type_utils.JitScalarType.FLOAT
         )
@@ -3886,7 +3886,8 @@ def new_zeros(
     g: jit_utils.GraphContext, self, sizes, dtype, layout, device, pin_memory=False
 ):
     self_dtype = symbolic_helper._try_get_scalar_type(self)
-    if dtype is None and self_dtype is not None:
+
+    if symbolic_helper._is_none(dtype) and self_dtype is not None:
         dtype = self_dtype
     return zeros(g, sizes, dtype, layout, device, pin_memory)
 
@@ -3929,7 +3930,7 @@ def ones_like(
     memory_format=None,
 ):
     shape = g.op("Shape", input)
-    if dtype is None:
+    if symbolic_helper._is_none(dtype):
         scalar_type = _type_utils.JitScalarType.from_value(
             input, _type_utils.JitScalarType.FLOAT
         )
@@ -3948,7 +3949,7 @@ def new_ones(
     g: jit_utils.GraphContext, self, sizes, dtype, layout, device, pin_memory=False
 ):
     self_dtype = symbolic_helper._try_get_scalar_type(self)
-    if dtype is None and self_dtype is not None:
+    if symbolic_helper._is_none(dtype) and self_dtype is not None:
         dtype = self_dtype
     return ones(g, sizes, dtype, layout, device, pin_memory)
 
@@ -4025,7 +4026,7 @@ def new_full(
     pin_memory=False,
 ):
     self_dtype = symbolic_helper._try_get_scalar_type(self)
-    if dtype is None and self_dtype is not None:
+    if symbolic_helper._is_none(dtype) and self_dtype is not None:
         dtype = self_dtype
     return full(g, size, fill_value, dtype, layout, device, pin_memory)
 


Test patch:
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
--- a/test/onnx/test_pytorch_onnx_onnxruntime.py
+++ b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -6511,6 +6511,21 @@ def forward(self, x):
         self.run_test(Zero_(), x, input_names=["x"], dynamic_axes={"x": [0, 1, 2]})
         self.run_test(Zero_(), x, remained_onnx_input_idx=[])
 
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_new_zeros_with_dtype(self):
+        class MyModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.emb = torch.nn.Embedding(50, 64)
+
+            def forward(self, x):
+                inp = x.new_zeros(x.shape)
+                return self.emb(inp)
+
+        model = MyModel()
+        x = torch.Tensor([[2, 5, 6], [3, 2, 5]]).to(torch.int64)
+        self.run_test(model, x, input_names=["x"], dynamic_axes={"x": [0, 1]})
+
     @skipIfUnsupportedMinOpsetVersion(9)
     def test_new_ones(self):
         class OnesModel(torch.nn.Module):