[ET-VK][ez] Restrict batch norm operator to 4D input tensors only

backends/vulkan/op_registry.py

@@ -793,17 +793,35 @@ def register_ported_op_all_packed_dims():
 
 
 # Ported ops that support their own prepacking.
-@update_features(
-    [
-        exir_ops.edge.aten.embedding.default,
-        exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
-    ]
-)
-def register_ported_ops_with_prepacking():
+@update_features(exir_ops.edge.aten.embedding.default)
+def register_embedding_op():
+    return OpFeatures(
+        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
+        supports_prepacking=True,
+        supports_resize=True,
+    )
+
+
+@update_features(exir_ops.edge.aten._native_batch_norm_legit_no_training.default)
+def register_batch_norm_op():
+    def check_batch_norm_node(node: torch.fx.Node) -> bool:
+        x = node.args[0]
+        if not isinstance(x, torch.fx.Node):
+            return False
+        x_val = x.meta.get("val", None)
+        if x_val is None:
+            return False
+        x_shape = x_val.size()
+        # Only support 4-D input tensors since this is a restriction enforced by the
+        # operator implementation.
+        # TODO(ssjia): Add shape agnostic support for batch norm
+        return len(x_shape) == 4
+
     return OpFeatures(
         inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
         supports_prepacking=True,
         supports_resize=True,
+        are_node_inputs_supported_fn=check_batch_norm_node,
     )
 
 

backends/vulkan/test/test_vulkan_delegate.py

@@ -1061,6 +1061,24 @@ def forward(self, x):
             sample_inputs,
         )
 
+    def test_vulkan_backend_batch_norm_after_linear(self):
+        class LinearBatchNormModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(128, 64)
+                self.bn = torch.nn.BatchNorm1d(num_features=64)
+
+            def forward(self, x):
+                x = self.linear(x)
+                return self.bn(x)
+
+        sample_inputs = (torch.randn(size=(4, 128), dtype=torch.float32),)
+
+        self.lower_module_and_test_output(
+            LinearBatchNormModule(),
+            sample_inputs,
+        )
+
     def test_vulkan_backend_full(self):
         class FullModule(torch.nn.Module):
             def __init__(self):