Vulkan linalg norm fails with arg out of range

🐛 Describe the bug

Vulkan models containing torch.linalg.norm can fail to partition with an index out of range error.

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower, EdgeCompileConfig, to_edge
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
from typing import Callable, List, Optional, Tuple, Union

class LinalgNormModel(torch.nn.Module):
    def __init__(
        self,
        ord: Optional[Union[int, float, str]] = None,
        dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None,
        keepdim: bool = False,
        dtype: Optional[torch.dtype] = None
    ):
        super().__init__()
        self.ord = ord
        self.dim = dim
        self.keepdim = keepdim
        self.dtype = dtype
        
    def forward(self, x):
        return torch.linalg.norm(x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype)
        
model = LinalgNormModel(ord=1)
inputs = (
    torch.randn(5, 10),
)
eager_outputs = model(*inputs)

ep = torch.export.export(model.eval(), inputs)
print(ep)
lowered = to_edge_transform_and_lower(
    ep,
    partitioner=[VulkanPartitioner()],
    compile_config=EdgeCompileConfig(_check_ir_validity=False)
).to_executorch()
print(lowered.exported_program())

et_model = _load_for_executorch_from_buffer(lowered.buffer)
et_outputs = et_model([*inputs])[0]

print(f"Inputs: {inputs}")
print(f"Eager: {eager_outputs}")
print(f"ET:    {et_outputs}")

Outputs:

...
File /data/users/gjcomer/fbsource/buck-out/v2/gen/fbcode/fdcb6705e87e1def/bento_kernels/cria/__bento_kernel_cria_binary__/bento_kernel_cria_binary#link-tree/executorch/backends/vulkan/op_registry.py:457, in register_reduce_op.<locals>.check_reduce_node(node)
    454 if isinstance(dim_list, list) and len(dim_list) != 1:
    455     return False
--> 457 keepdim = node.args[2]
    458 if isinstance(keepdim, bool) and not keepdim:
    459     return False
IndexError: tuple index out of range

Versions

Run on Meta internal master, Jul 3, fbcode/SwiftShader

cc @SS-JIA @manuelcandales @cbilgin

Start in executorch/backends/vulkan/op_registry.py at register_reduce_op and its check_reduce_node helper, especially the access to node.args[2]. Reproduce the failure with the LinalgNormModel example using ord=1, then verify that Vulkan partitioning handles this exported torch.linalg.norm form without an index-out-of-range error.