This document is relevant for: Inf2, Trn1, Trn2

PyTorch Neuron (torch-neuronx) - Supported Operators

Operator support

The following list the aten operators supported by torch-neuronx.

aten::_s_where

aten::_softmax

aten::_softmax_backward_data

aten::_unsafe_view

aten::add

aten::addcdiv_

aten::addcmul

aten::addmm

aten::bernoulli_

aten::bmm

aten::constant_pad_nd

aten::div

aten::embedding

aten::embedding_dense_backward

aten::empty

aten::expand

aten::fill_

aten::index_select

aten::_log_softmax

aten::_log_softmax_backward_data

aten::lt

aten::mm

aten::mul

aten::native_batch_norm

aten::native_batch_norm_backward

aten::neg

aten::permute

aten::relu

aten::rsub

aten::select

aten::slice

aten::sqrt

aten::sum

aten::t

aten::tanh

aten::tanh_backward

aten::threshold_backward

aten::transpose

aten::unsqueeze

aten::view

aten::zero_

This document is relevant for: Inf2, Trn1, Trn2