unsupported by backend lowering: tensor with unknown rank or dtype” errors when modifying module attributes

When we try to convert our pytorch model using torch-mlir, we ran into several “unsupported by backend lowering: tensor with unknown rank or dtype” errors during the torch-backend-to-linalg-on-tensors-backend-pipeline pass. We notice that it happens when we try to modify attributes that are initialized in __init__ function.

To reproduce this error, we wrote a simple test that is basically a single identical layer, but we modified its attributes in the forward function.

```
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder

from torch_mlir.passmanager import PassManager
from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend
from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export

mb = ModuleBuilder()

class IdenticalModule(torch.nn.Module):
def __init__(self):
super(IdenticalModule, self).__init__()
self.attr0:Tensor = torch.tensor(1.)
def helper(self):
self.attr0 *= 2.
pass
def forward(self, x:Tensor) -> Tensor:
self.helper()
return x

class TestModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.s = IdenticalModule()
self.train(False)

def forward(self, x) -> Tensor:
return self.s.forward(x)

test_module = TestModule()
scripted_cell = torch.jit.script(test_module)
scripted_cell.save('scripted_cell.pt')
class_annotator = ClassAnnotator()

class_annotator.exportNone(scripted_cell._c._type())
class_annotator.exportPath(scripted_cell._c._type(), ["forward"])
class_annotator.annotateArgs(
scripted_cell._c._type(),
["forward"],
[None, ([-1, -1], torch.float32, True)],
)

mb.import_module(scripted_cell._c, class_annotator)
print("-------------------------------init IR-------------------------------")
mb.module.dump()

backend = refbackend.RefBackendLinalgOnTensorsBackend()
with mb.module.context:
pm = PassManager.parse('torchscript-module-to-torch-backend-pipeline')
print("-----------------------torch-backend------------------")
pm.run(mb.module)
mb.module.dump()
pm = PassManager.parse('torch-backend-to-linalg-on-tensors-backend-pipeline')
print("-----------------------linalg-------------------------")
pm.run(mb.module)
mb.module.dump()

compiled = backend.compile(mb.module)
jit_module = backend.load(compiled)

def predictions(torch_func, jit_func, input1):
print("PyTorch prediction")
torch_prediction = torch_func(input1)
print(torch_prediction)
print("MLIR prediction")
jit_prediction = jit_func(input1.numpy())
print(jit_prediction)

x = torch.rand(4, 3)
predictions(test_module.forward, jit_module.forward, x)

```

We printed out the mlir codes and found that error occurs on “torch.copy.to_vtensor” when copying attributes (`%1 = torch.copy.to_vtensor %0 : !torch.vtensor`):

```
-------------------------------init IR-------------------------------
module attributes {torch.debug_module_name = "TestModule"} {
func private @__torch__.IdenticalModule.forward(%arg0: !torch.nn.Module<"__torch__.IdenticalModule">, %arg1: !torch.tensor) -> !torch.tensor {
%3 = torch.prim.CallMethod %arg0["helper"] () : !torch.nn.Module<"__torch__.IdenticalModule">, () -> !torch.none
return %arg1 : !torch.tensor
}
func private @__torch__.IdenticalModule.helper(%arg0: !torch.nn.Module<"__torch__.IdenticalModule">) -> !torch.none {
%none_0 = torch.constant.none
%float2.000000e00 = torch.constant.float 2.000000e+00
%3 = torch.prim.GetAttr %arg0["attr0"] : !torch.nn.Module<"__torch__.IdenticalModule"> -> !torch.tensor
%4 = torch.aten.mul_.Scalar %3, %float2.000000e00 : !torch.tensor, !torch.float -> !torch.tensor
torch.prim.SetAttr %arg0["attr0"] = %4 : !torch.nn.Module<"__torch__.IdenticalModule">, !torch.tensor
return %none_0 : !torch.none
}
func private @__torch__.TestModule.forward(%arg0: !torch.nn.Module<"__torch__.TestModule">, %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[?,?],f32>}) -> !torch.tensor {
%3 = torch.prim.GetAttr %arg0["s"] : !torch.nn.Module<"__torch__.TestModule"> -> !torch.nn.Module<"__torch__.IdenticalModule">
%4 = torch.prim.CallMethod %3["forward"] (%arg1) : !torch.nn.Module<"__torch__.IdenticalModule">, (!torch.tensor) -> !torch.tensor
return %4 : !torch.tensor
}
torch.class_type @__torch__.TestModule {
torch.attr private "training" : !torch.bool
torch.attr private "_is_full_backward_hook" : !torch.optional
torch.attr private "s" : !torch.nn.Module<"__torch__.IdenticalModule">
torch.method "forward", @__torch__.TestModule.forward
}
%false = torch.constant.bool false
%none = torch.constant.none
torch.class_type @__torch__.IdenticalModule {
torch.attr private "training" : !torch.bool
torch.attr private "_is_full_backward_hook" : !torch.optional
torch.attr private "attr0" : !torch.tensor
torch.method private "forward", @__torch__.IdenticalModule.forward
torch.method private "helper", @__torch__.IdenticalModule.helper
}
%0 = torch.tensor.literal(dense<2.000000e+00> : tensor) : !torch.tensor<[],f32>
%1 = torch.nn_module {
torch.slot "training", %false : !torch.bool
torch.slot "_is_full_backward_hook", %none : !torch.none
torch.slot "attr0", %0 : !torch.tensor<[],f32>
} : !torch.nn.Module<"__torch__.IdenticalModule">
%2 = torch.nn_module {
torch.slot "training", %false : !torch.bool
torch.slot "_is_full_backward_hook", %none : !torch.none
torch.slot "s", %1 : !torch.nn.Module<"__torch__.IdenticalModule">
} : !torch.nn.Module<"__torch__.TestModule">
}
----------------------------------torch-backend-----------------------------------
module attributes {torch.debug_module_name = "TestModule"} {
torch.global_slot "private" @s.attr0 : !torch.tensor {
%0 = torch.tensor.literal(dense<2.000000e+00> : tensor) : !torch.tensor<[],f32>
torch.global_slot.init %0 : !torch.tensor<[],f32>
}
func @forward(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
%float2.000000e00 = torch.constant.float 2.000000e+00
%0 = torch.global_slot.get @s.attr0 : !torch.tensor

%1 = torch.copy.to_vtensor %0 : !torch.vtensor <-- Seems like rank and dtype cannot be detected when copying attributes

%2 = torch.aten.mul.Scalar %1, %float2.000000e00 : !torch.vtensor, !torch.float -> !torch.vtensor
torch.overwrite.tensor %2 overwrites %0 : !torch.vtensor, !torch.tensor
torch.global_slot.set @s.attr0 = %0 : !torch.tensor
return %arg0 : !torch.vtensor<[?,?],f32>
}
}
----------------------------------linalg-----------------------------------
error: unsupported by backend lowering: tensor with unknown rank or dtype
note: see current operation: %2 = "torch.copy.to_vtensor"(%1) : (!torch.tensor) -> !torch.vtensor
note: this is likely due to a missing case in RefineTypes
Traceback (most recent call last):
File "/work/shared/users/ugrad/zh338/torch-mlir/examples/simple_test.py", line 69, in
pm.run(mb.module)
RuntimeError: Failure while executing pass pipeline.

```

Start with examples/simple_test.py and the torch-backend-to-linalg-on-tensors-backend-pipeline shown in the traceback. Inspect the RefineTypes handling for torch.copy.to_vtensor after torch.global_slot.get produces the modified module attribute. Done means the reproduced model passes this lowering pipeline without the unknown-rank-or-dtype error and can continue to compilation.

Written by the indexing model from the issue text.