Update on "[ET-VK][ez] Add back tensor dim check"

## Context

Vulkan cannot represent higher dimensional tensors (tensors with dim > 4) at the moment, but due to some refactors implemented last year the partitioner check to avoid lowering ops that involve high dimensional tensors was accidentally removed.

This diff adds back the check, as well as a test to verify that high dimensional tensors do not get lowered.

Differential Revision: [D68630966](https://our.internmc.facebook.com/intern/diff/D68630966/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -84,7 +84,7 @@ def op_node_is_compatible(
             features = get_op_features(target)
 
         # Check for high dimensional tensors
-        if utils.tensor_node_is_high_dim(node):
+        if utils.is_tensor_node(node) and utils.tensor_node_is_high_dim(node):
             return False, "contains high dim tensor"
 
         valid_texture_layouts = utils.possible_node_memory_layouts(
@@ -99,7 +99,7 @@ def op_node_is_compatible(
                 and i not in features.skip_limits_check
             ):
                 # Check for high dimensional tensors
-                if utils.tensor_node_is_high_dim(arg):
+                if utils.is_tensor_node(arg) and utils.tensor_node_is_high_dim(arg):
                     return False, "contains high dim tensor"
 
                 arg_texture_layouts = utils.possible_node_memory_layouts(

backends/vulkan/utils.py

@@ -132,20 +132,16 @@ def within_buffer_limit(node: torch.fx.Node, buffer_limit: int) -> int:
 
 def tensor_node_is_high_dim(node: torch.fx.Node) -> bool:
     """
-    If the node does not contain a tensor or a collection of tensors, return False.
-    Otherwise, return True if the tensor is high dimensional (i.e. rank > 4).
-    """
-    if is_tensor_node(node):
-        if isinstance(node.meta["val"], FakeTensor):
-            return len(node.meta["val"].shape) > 4
-        if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
-            for fake_tensor in node.meta["val"]:
-                if isinstance(fake_tensor, FakeTensor):
-                    if len(fake_tensor.shape) > 4:
-                        return True
-        return False
-    else:
-        return False
+    Returns true if a given node contains a tensor with more than 4 dimensions
+    """
+    if isinstance(node.meta["val"], FakeTensor):
+        return len(node.meta["val"].shape) > 4
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        for fake_tensor in node.meta["val"]:
+            if isinstance(fake_tensor, FakeTensor):
+                if len(fake_tensor.shape) > 4:
+                    return True
+    return False
 
 
 def required_image_extents(sizes: torch.Size, layout: VkMemoryLayout) -> ImageExtents: