Repository: apache/tvm

Issue:
[Bug] [Relay][PyTorch] `aten::new_zeros` support is broken with torch>=2.1.0
Seems like the [aten::new_zeros](https://pytorch.org/docs/stable/generated/torch.Tensor.new_zeros.html) support is broken with torch>=2.1.0. #14747 

### Steps to reproduce

```python
from tvm import relay
import torch

torch.set_grad_enabled(False)

class NewZeros1(torch.nn.Module):
  def forward(self, x):
    return x.new_zeros((2, 3))

input_data = torch.tensor((), dtype=torch.float)
module = NewZeros1().float().eval()
scripted_module = torch.jit.trace(module, input_data)
input_infos = [('x', (input_data.shape, 'float32'))]
mod, params = relay.frontend.from_pytorch(scripted_module, input_infos)
```

### Actual behavior

```
Traceback (most recent call last):
  File "/home/ubuntu/workspace/sandbox/tvm_/frontend/new_zeros.py", line 17, in <module>
    mod, params = relay.frontend.from_pytorch(scripted_module, input_infos)
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/relay/frontend/pytorch.py", line 5183, in from_pytorch
    outputs = converter.convert_operators(operator_nodes, outputs, ret_name)
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/relay/frontend/pytorch.py", line 4402, in convert_operators
    relay_out = relay_op(
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/relay/frontend/pytorch.py", line 851, in new_zeros
    return self.full_impl(data, 0, dtype)
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/relay/frontend/pytorch.py", line 755, in full_impl
    out = _op.full(fill_value, size, dtype=dtype)
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/relay/op/transform.py", line 535, in full
    return _make.full(fill_value, shape, dtype)
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 239, in __call__
    raise_last_ffi_error()
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error
    raise py_err
tvm._ffi.base.TVMError: Traceback (most recent call last):
  2: _ZN3tvm7runtime13PackedFuncObj
  1: tvm::runtime::TypedPackedFunc<tvm::RelayExpr (tvm::RelayExpr, tvm::runtime::Array<tvm::Integer, void>, tvm::runtime::DataType)>::AssignTypedLambda<tvm::RelayExpr (*)(tvm::RelayExpr, tvm::runtime::Array<tvm::Integer, void>, tvm::runtime::DataType)>(tvm::RelayExpr (*)(tvm::RelayExpr, tvm::runtime::Array<tvm::Integer, void>, tvm::runtime::DataType), std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const
  0: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::runtime::DataType<tvm::runtime::DataType>() const
  3: _ZN3tvm7runtime13PackedFuncObj
  2: tvm::runtime::TypedPackedFunc<tvm::RelayExpr (tvm::RelayExpr, tvm::runtime::Array<tvm::Integer, void>, tvm::runtime::DataType)>::AssignTypedLambda<tvm::RelayExpr (*)(tvm::RelayExpr, tvm::runtime::Array<tvm::Integer, void>, tvm::runtime::DataType)>(tvm::RelayExpr (*)(tvm::RelayExpr, tvm::runtime::Array<tvm::Integer, void>, tvm::runtime::DataType), std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const
  1: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::runtime::DataType<tvm::runtime::DataType>() const
  0: tvm::runtime::TVMArgValue::operator DLDataType() const
  File "/home/ubuntu/workspace/sandbox/.dep/tvm/include/tvm/runtime/packed_func.h", line 777
TVMError: In function relay.op._make.full(0: RelayExpr, 1: Array<IntImm>, 2: DataType) -> RelayExpr: error while converting argument 2: [19:47:16] /home/ubuntu/workspace/sandbox/.dep/tvm/include/tvm/runtime/packed_func.h:2210: InternalError: Check failed: type_code_ == kTVMDataType (8 vs. 5) : expected DLDataType but got Object
```

### Environment

TVM: c8ef902a752e71dc92fb7fdfd7c3fcc4178a3455
LLVM=ON
CUDA=OFF
Torch==2.1.0

### Triage

Please refer to the list of label tags [here](https://github.com/apache/tvm/wiki/Issue-Triage-Labels) to find the relevant tags and add them below in a bullet format (example below).

* frontend:pytorch


cc @shingjan @yelite


Discussion:

Resolving patch:
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -847,7 +847,7 @@ def new_zeros(self, inputs, input_types):
             dtype = _convert_dtype_value(inputs[2])
         else:
             # if dtype is None, use the dtype of the input tensor
-            dtype = self.infer_type(inputs[0])
+            dtype = self.infer_type(inputs[0]).dtype
         return self.full_impl(data, 0, dtype)
 
     def full(self, inputs, input_types):
@@ -898,7 +898,7 @@ def new_full(self, inputs, input_types):
             dtype = _convert_dtype_value(inputs[3])
         else:
             # if dtype is None, use the dtype of the input tensor
-            dtype = self.infer_type(inputs[0])
+            dtype = self.infer_type(inputs[0]).dtype
 
         return self.full_impl(data, fill_value, dtype)
 


Test patch: