Handle complex element type in torch.vtensor conversion (#175)

Signed-off-by: Boian Petkantchin <boian.petkantchin@amd.com>
iree-org · Sep 27, 2024 · 3410f23 · 3410f23
1 parent a1da408
commit 3410f23
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/shark_turbine/dynamo/type_conversion.py b/shark_turbine/dynamo/type_conversion.py
@@ -32,7 +32,7 @@
 #   1. Local name (int, float, vtensor)
 #   2. Parameter block ("<...>"), including the delimitters
 #   3. Inner parameter block (no delimitters)
-DECOMPOSE_TORCH_TYPE_PATTERN = re.compile(r"^!torch.([^<]+)(<([^>]*)>)?$")
+DECOMPOSE_TORCH_TYPE_PATTERN = re.compile(r"^!torch\.([^<]+)(<(.*)>)?$")
 
 # Decomposes a vtensor parameter block into a dimension list and dtype. Groups:
 #   1. Dimension list

diff --git a/tests/dynamo/type_conversion_test.py b/tests/dynamo/type_conversion_test.py
@@ -32,6 +32,7 @@ def testValueTensors(self):
         self._compareNative("!torch.vtensor<[2, 2],f32>", "tensor<2x2xf32>")
         self._compareNative("!torch.vtensor<[?, ?],f32>", "tensor<?x?xf32>")
         self._compareNative("!torch.vtensor<[],f32>", "tensor<f32>")
+        self._compareNative("!torch.vtensor<[],complex<f32>>", "tensor<complex<f32>>")
 
     def _compareNative(self, torch_str: str, native_str: str, *, signless: bool = True):
         with self.conv._context: