1# Copyright 2023 Arm Limited and/or its affiliates. 2# 3# This source code is licensed under the BSD-style license found in the 4# LICENSE file in the root directory of this source tree. 5 6# pyre-unsafe 7from typing import List 8 9import serializer.tosa_serializer as ts 10import torch 11from executorch.backends.arm.operators.node_visitor import ( 12 NodeVisitor, 13 register_node_visitor, 14) 15from executorch.backends.arm.tosa_mapping import TosaArg 16from serializer.tosa_serializer import TosaOp 17 18 19@register_node_visitor 20class GetItemVisitor(NodeVisitor): 21 target = "getitem" 22 23 def __init__(self, *args): 24 super().__init__(*args) 25 26 def define_node( 27 self, 28 node: torch.fx.Node, 29 tosa_graph: ts.TosaSerializer, 30 inputs: List[TosaArg], 31 output: TosaArg, 32 is_quant_node: bool, 33 ) -> None: 34 item_name = inputs[0].name 35 ## Simply add an identityOp 36 tosa_graph.addOperator(TosaOp.Op().IDENTITY, [item_name], [output.name]) 37