diff --git a/python/tvm/relay/backend/contrib/ethosu/codegen.py b/python/tvm/relay/backend/contrib/ethosu/codegen.py index 1f331822e1ac..e51f1702773b 100644 --- a/python/tvm/relay/backend/contrib/ethosu/codegen.py +++ b/python/tvm/relay/backend/contrib/ethosu/codegen.py @@ -90,14 +90,14 @@ def visit_call(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Call: new_call = call lut_activations = ["TANH", "LUT"] - if ( - call.op.name == "contrib.ethosu.identity" - and call.attrs.activation in lut_activations - and isinstance(call.args[0], tvm.relay.expr.Call) - ): + if isinstance(call.op, tvm.ir.Op) and isinstance(call.args[0], tvm.relay.expr.Call): producer_op = call.args[0] # Check if the producer can do a LUT operation - if producer_op.op.name in self.lut_ops.keys(): + if ( + producer_op.op.name in self.lut_ops.keys() + and call.op.name == "contrib.ethosu.identity" + and call.attrs.activation in lut_activations + ): # Check the producer doesn't already have a LUT has_lut = producer_op.attrs.activation in lut_activations if not has_lut: diff --git a/python/tvm/relay/backend/contrib/ethosu/op/op_attrs.py b/python/tvm/relay/backend/contrib/ethosu/op/op_attrs.py index e38a3dfd97de..a52736fe3964 100644 --- a/python/tvm/relay/backend/contrib/ethosu/op/op_attrs.py +++ b/python/tvm/relay/backend/contrib/ethosu/op/op_attrs.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""The attributes node used for EthosU Relay operators.""" +"""The attributes node used for Arm(R) Ethos(TM)-U NPU Relay operators.""" from tvm.ir import Attrs import tvm._ffi