ATen ops list

[justinchuby]

Canonical ops: https://docs-preview.pytorch.org/90644/ir.html

The list below is an incomplete list of canonical ops + other ops ordered by usage. canonical ops are marked with canonical. Be sure the check the page above for the correct signature for the canonical version of an op.

OP	GPT2	Type	PR
- aten::empty_strided			#432
- aten::empty			#264
- aten::to			what's this?
- aten::copy_			#432
- aten::as_strided			something like pool
- aten::fill_			#432
- aten::_to_copy		canonical	what's this?
- aten::empty_like			#264
- aten::zero_			not defined in core.py
- aten::detach	Yes		#324
- aten::zeros			#251
- aten::detach_			not defined in core.py
- aten::view	Yes	canonical	#265
- aten::uniform_			not defined in core.py
- aten::add_			not defined in core.py
- aten::_has_compatible_shallow_copy_type			not defined in core.py
- aten::_reshape_alias			not defined in core.py
- aten::reshape			#265
- aten::add	Yes	canonical	#235
- aten::sum	Yes	canonical	#348
- aten::transpose	Yes		done
- aten::resize_			not defined in core.py
- aten::expand	Yes	canonical	#265
- aten::ones_like			#247
- aten::t			#252
- aten::mul	Yes	canonical	#235
- aten::linear			done
- aten::mm		canonical	#252
- aten::new_empty_strided			#447
- aten::clone	Yes	canonical	#265
- aten::_local_scalar_dense			not defined in core.py
- aten::item			is there any op return scalar?
- aten::slice	Yes	canonical	#265
- aten::div	Yes	canonical	#265
- aten::normal_			#440
- aten::_log_softmax		canonical	no definition
- aten::log_softmax		canonical	#264
- aten::unsqueeze	Yes	canonical	#266
- aten::select			done
- aten::cat		canonical	#265
- aten::eq		canonical	#265
- aten::_cat			no definition
- aten::narrow			#440
- aten::_log_softmax_backward_data			backward
- aten::zeros_like			#251
- aten::ones			done
- aten::cross_entropy_loss			#444
- aten::masked_fill_			done
- aten::stack			not defined
- aten::nll_loss_nd			#453
- aten::_unsafe_view			no define
- aten::contiguous			#454
- aten::sub	Yes	canonical	#235
- aten::dropout			#454
- aten::mean		canonical	xiaowu
- aten::nll_loss_forward			forward
- aten::nll_loss			#453
- aten::matmul			#252
- aten::resolve_conj			#458
- aten::set_			no define
- aten::permute	Yes	canonical	#293
- aten::nll_loss_backward			backward
- aten::mul_			done
- aten::bmm	Yes	canonical	#255
- aten::resolve_neg			#458
- aten::layer_norm			#459
- aten::_softmax		canonical	no define
- aten::softmax			#313
- aten::native_layer_norm	Yes	canonical	#330
- aten::clamp_min			#247
- aten::squeeze			#460
- aten::new_empty			#447
- aten::native_layer_norm_backward			backward
- aten::_softmax_backward_data			backward
- aten::arange	Yes		#285
- aten::native_dropout		canonical	#432
- aten::native_dropout_backward			backward
- aten::index_select		canonical	#274
- aten::addmm	Yes	canonical	#255
- aten::neg		canonical	#266
- aten::flatten			no define
- aten::ne		canonical	#266
- aten::abs		canonical	#255
- aten::embedding	Yes		#306
- aten::unbind			no define
- aten::any			#463
- aten::expand_as			#463
- aten::embedding_backward			backward
- aten::embedding_dense_backward		canonical	backward
- aten::is_nonzero			#530
- aten::split	Yes		#337
- aten::le		canonical	#325
- aten::select_backward			backward
- aten::max		canonical	#314
- aten::random_			no def
- aten::rsub	Yes		#279
- aten::iand			no def
- aten::chunk			#511
- aten::slice_backward			backward
- aten::split_with_sizes			#513
- aten::bitwise_and_			#266
- aten::isfinite			#513
- aten::flatten_dense_tensors
- aten::all			#513
- aten::_index_put			#522
- aten::unflatten_dense_tensors
- aten::index_put_			#522
- aten::isclose			#526
- aten::ior
- aten::bitwise_or_			#266
- aten::alias			done
- aten::allclose			#526
- aten::gelu		canonical	#323
- aten::gelu_backward			backward
- aten::isinf		canonical	#266
- aten::gt		canonical	#247
- aten::isnan			#530
- aten::sub_			done
- aten::round			#247
- aten::unfold			#534
- aten::repeat		canonical	#247
- aten::lt		canonical	#247
- aten::masked_fill			done
- aten::argmax		canonical	#324
- aten::argmin		canonical	#324
- aten::clamp		canonical	#247
- aten::sigmoid		canonical	#266
- aten::convolution		canonical	#411
- aten::convolution_backward		canonical	backward
- aten::index
- aten::new_zeros			#535
- aten::cudnn_convolution
- aten::sqrt		canonical	#266
- aten::result_type
- aten::threshold_backward			backward
- aten::where	Yes		#286
- aten::_s_where
- aten::ge		canonical	#325
- aten::broadcast_tensors
- aten::mse_loss_backward			backward
- aten::mse_loss			#536
- aten::binary_cross_entropy_with_logits_backward			backward
- aten::addcmul_
- aten::binary_cross_entropy_with_logits
- aten::conv2d
- aten::addcdiv_
- aten::cumsum
- aten::as_strided_
- aten::constant_pad_nd		canonical	#492
- aten::nonzero		canonical	#266
- aten::tanh	Yes	canonical	#256
- aten::scatter_
- aten::randint
- aten::min		canonical	#491
- aten::full	Yes	canonical	#286
- aten::pow	Yes	canonical	#293
- aten::bitwise_and		canonical	#266
- aten::einsum
- aten::scalar_tensor		canonical	#492
- aten::norm
- aten::relu		canonical	#266
- aten::remainder		canonical	#266
- aten::gather
- aten::tanh_backward			backward
- aten::repeat_interleave
- aten::topk
- aten::randperm
- aten::reciprocal		canonical	#266
- aten::new_ones			#535
- aten::type_as
- aten::sort
- aten::exponential_
- aten::multinomial			#1032
- aten::index_add_
- aten::scatter
- aten::bincount
- aten::take
- aten::randn
- aten::cos		canonical	#256
- aten::sin		canonical	#256
- aten::exp	Yes	canonical	#256
- aten::tril
- aten::_batch_norm_impl_index
- aten::unsqueeze_
- aten::batch_norm
- aten::cudnn_batch_norm
- aten::relu_
- aten::squeeze_
- aten::max_pool2d_with_indices		canonical	Jay
- aten::cudnn_batch_norm_backward			backward
- aten::max_pool2d
- aten::dense_dim
- aten::to_dense
- aten::coalesced
- aten::sparse_dim
- aten::copy_sparse_to_sparse_
- aten::_nnz
- aten::is_coalesced
- aten::_indices
- aten::sparse_coo_tensor
- aten::_values
- aten::_sparse_coo_tensor_with_dims_and_tensors
- aten::full_like			#286
- aten::log		canonical	#281
- aten::max_pool2d_with_indices_backward		canonical	backward
- aten::bitwise_not		canonical	#266
- aten::adaptive_avg_pool2d
- aten::rand
- aten::bernoulli_
- aten::divide_
- aten::conv1d			#431
- aten::conv3d			#431
- aten::sigmoid_backward			backward
- aten::and
- aten::minimum		canonical	#313
- aten::gather_backward			backward
- aten::linspace
- aten::scatter_add_
- aten::masked_select
- aten::flip
- aten::_conv_depthwise2d
- aten::index_put			#522
- aten::upsample_bilinear2d		canonical	seems deprecated
- aten::native_group_norm		canonical	no torch def, and no def in torch-onnx symbolic function
- aten::group_norm
- aten::native_group_norm_backward		canonical	backward
- aten::upsample_bilinear2d_backward			backward
- aten::one_hot
- aten::resize_as_
- aten::_weight_norm_cuda_interface
- aten::_weight_norm
- aten::norm_except_dim
- aten::bitwise_or		canonical	#266
- aten::_use_cudnn_ctc_loss
- aten::_ctc_loss
- aten::nll_loss2d			#453
- aten::index_select_backward			backward
- aten::ctc_loss
- aten::kl_div_backward			backward
- aten::kl_div
- aten::or
- aten::_ctc_loss_backward			backward
- aten::triu
- aten::value_selecting_reduction_backward
- aten::nll_loss2d_backward			backward
- aten::nll_loss2d_forward
- aten::view_as
- aten::is_pinned
- aten::avg_pool2d_backward		canonical	backward
- aten::avg_pool2d		canonical	Jay
- aten::_weight_norm_cuda_interface_backward			backward
- aten::rand_like
- aten::numpy_T
- aten::std
- aten::rsqrt		canonical	#266
- aten::conv_transpose2d
- aten::cudnn_convolution_transpose
- aten::l1_loss
- aten::im2col
- aten::sgn_
- aten::roll
- aten::eye
- aten::glu
- aten::l1_loss_backward			backward
- aten::nonzero_numpy
- aten::im2col_backward			backward
- aten::_adaptive_avg_pool2d		canonical	Jay
- aten::maximum		canonical	#313
- aten::ceil
- aten::unique_consecutive
- aten::argsort
- aten::is_same_size			#530
- aten::equal			#265
- aten::erf		canonical	#266
- aten::mkldnn_convolution
- aten::_adaptive_avg_pool2d_backward		canonical	backward
- aten::meshgrid
- aten::floor_
- aten::bernoulli
- aten::amax	Yes	canonical	#279
- aten::pixel_shuffle
- aten::logical_not		canonical	#266
- aten::avg_pool1d			Jay
- aten::logsumexp			#281
- aten::leaky_relu_backward			backward
- aten::lift
- aten::new_full			#286
- aten::upsample_nearest2d		canonical	#279
- aten::max_pool1d
- aten::frobenius_norm
- aten::triu_
- aten::feature_dropout
- aten::square
- aten::nansum
- aten::transpose_
- aten::sqrt_
- aten::cosine_similarity
- aten::view_as_real
- aten::real
- aten::dropout_
- aten::upsample_nearest2d_backward			backward
- aten::adaptive_avg_pool1d
- aten::clip
- aten::sgn
- aten::abs_
- aten::leaky_relu		canonical	#266
- aten::silu_backward			backward
- aten::fmod		canonical	#266
- aten::broadcast_to
- aten::erfinv_
- aten::logical_and		canonical	#266
- aten::hardsigmoid
- aten::hardswish_
- aten::hann_window
- aten::fill_diagonal_
- aten::t_
- aten::isneginf			#530
- aten::from_file
- aten::smooth_l1_loss
- aten::grid_sampler_2d_backward			backward
- aten::replication_pad1d
- aten::amp_update_scale
- aten::grid_sampler_2d
- aten::grid_sampler
- aten::amp_foreach_non_finite_check_and_unscale
- aten::dot
- aten::silu_
- aten::_convolution_double_backward
- aten::instance_norm
- aten::pin_memory
- aten::reshape_as
- aten::_unique2
- aten::baddbmm
- aten::index_fill_
- aten::index_fill
- aten::outer
- aten::ger
- aten::index_copy
- aten::index_copy_
- aten::_pin_memory
- aten::_fft_r2c
- aten::fft_rfft
- aten::elu_backward			backward
- aten::elu
- aten::selu
- aten::movedim
- aten::moveaxis
- aten::median
- aten::glu_backward			backward
- aten::_cudnn_rnn_backward			backward
- aten::_cudnn_rnn
- aten::cudnn_is_acceptable
- aten::lstm
- aten::_cudnn_rnn_flatten_weight
- aten::_use_cudnn_rnn_flatten_weight
- aten::reflection_pad2d
- aten::silu
- aten::randn_like
- aten::cumprod
- aten::masked_scatter_
- aten::masked_select_backward			backward
- aten::native_batch_norm_backward			backward
- aten::hardtanh_backward			backward
- aten::fake_quantize_per_tensor_affine_cachemask_backward			backward
- aten::hardtanh_
- aten::native_batch_norm		canonical	when training=True, the return mean,var is not close with torch test case
- aten::fake_quantize_per_channel_affine_cachemask
- aten::_fake_quantize_per_tensor_affine_cachemask_tensor_qparams
- aten::aminmax
- aten::_fused_moving_avg_obs_fq_helper
- aten::fused_moving_avg_obs_fake_quant
- aten::max_pool1d_with_indices
- aten::view_as_complex
- aten::_fft_c2c
- aten::fft_fftn
- aten::bitwise_xor		canonical	#266
- aten::log2			#281
- aten::sign		canonical	#266
- aten::floor
- aten::range
- aten::logical_and_			#266
- aten::logical_not_			#266
- aten::tile
- aten::diag_backward			backward
- aten::acos		canonical	#256
- aten::unique_dim
- aten::diag
- aten::_cdist_forward
- aten::cdist
- aten::hardsigmoid_backward			backward
- aten::hardswish_backward			backward
- aten::prod
- aten::binary_cross_entropy_backward			backward
- aten::binary_cross_entropy
- aten::leaky_relu_			#266
- aten::var_mean
- aten::_thnn_fused_gru_cell_backward			backward
- aten::_thnn_fused_gru_cell
- aten::gru_cell
- aten::multiply
- aten::softplus_backward			backward
- aten::polygamma
- aten::digamma
- aten::lgamma
- aten::softplus

### Tasks

[ramkrishna2910]

Some of the ops above are not a part of the ops_db in torch (which is a requirement for aten support).
For those ops do you recommend concurrent PRs to both onnx-script and pytorch? @justinchuby

[justinchuby]

I think a hand constructed test in onnxscript (OpInfo style if doable - you can simply create an OpInfo and append it to the list, normal function calls otherwise) will do. A PR to PyTorch is certainly welcome but that will also depend on what the ops_db maintainers think.

[zhiqwang]

Is there any way we can acquire an ATen IR graph?

[justinchuby]

@zhiqwang Can you elaborate?

[justinchuby]

@xiaowuhu @fatcat-z some ops I have seen recently: torch vision ops in https://github.com/pytorch/vision/blob/main/torchvision/ops/_register_onnx_ops.py (pending validation since we don't know what dynamo will give us). unflatten (#599) (very useful for recent transformer models), diag, diagonal

[ramkrishna2910]

@justinchuby in some cases there are ops coming from extensions of torch such as torch_scatter that are not supported by ONNX and they are not a part of aten ops either. (E.g. scatter_max, scatter_min).
Is the scope of this effort only to support aten ops or can we also include non-aten ops like the above?

[justinchuby]

We plan to support aten ops only but it should be very easy to create other ops as needed when one sees them

[titaiwangms]

cc @thiagocrepaldi

[justinchuby]

New ops from https://github.com/microsoft/onnx-converters-private/issues/166:

• aten::unbind.int Add op (unbind) | feat(torchlib) #831
• aten::index.Tensor Implement aten::index | feat(torchlib) (#862) #883
  • https://github.com/pytorch/pytorch/blob/f61b248d5b0c54db064fc27826c90d5956e0dfc0/aten/src/ATen/native/TensorAdvancedIndexing.cpp
• aten::unfold Add op(unfold) | feat(torchlib) #893
• prims::device_put Make aten::contiguous and device_put no-op | fix(torchlib) #835
• aten::linspace Add op (linspace) | feat(torchlib) #838
• aten::_embedding_bag AddOp(embedding bag) | feat(torchlib) #909
• aten::copy_ FX should handle this [ONNX] FX exporter: replace aten::copy_ with out-place version pytorch/pytorch#103857
• aten::_fft_c2c Implement _fft_* ops | feat(torchlib) #926
• aten::view_as_real: no-op Add ops (view*) | feat(torchlib) #814
• aten::_unsafe_index.Tensor Implement aten::index | feat(torchlib) (#862) #883
• aten::addcmul Implement add* ops | feat(torchlib) (#785) #791
• aten::roll Add op(roll) | feat(torchlib) #905
• aten::linalg_vector_norm AddOp(linalg_vector_norm) | feat(torchlib) #908
• aten::select_scatter.default Add op(select_scatter) | feat(torchlib) #901

[BowenBao]

prims::device_put FX should handle this as an no-op

why not implement in torchlib as no-op?

[justinchuby]

ONNX Script doesn't have the concept of a device and we generally filter everything device-related out. I think it would be good if device info are stripped away before translation.

[BowenBao]

Do you have context where this filter everything device-related out takes place?

[justinchuby]

We ignore the device arg here: https://github.com/pytorch/pytorch/blob/d3ba8901d8640eb16f88b2bfef9df7fa383d4b47/torch/onnx/_internal/fx/fx_onnx_interpreter.py#L171C1-L171C1

[justinchuby]

From https://github.com/microsoft/onnx-converters-private/issues/166#issuecomment-1658943209

• [torchlib] _embedding_bag #1021
• [torchlib] Implement aten::_softmax to avoid decomposition in exporter #857
• aten::multinomial Add Op(multinomial) | torchlib(feat) #1032
• aten::_native_multi_head_attention (from torch doc, not supported yet)
• aten::_scaled_dot_product_flash_attention Add Op(_scaled_dot_product_flash_attention) | feat(torchlib) #1043
• l1_loss
• [torchlib] Bug with aten::slice_scatter #945
• varmean: Shape error (check_model cannot pass, but inference can pass. May be related with 2 outputs for this function)
• [torchlib] Implement diagonal and diag #679
• https://github.com/microsoft/onnxscript/issues?q=is%3Aopen+is%3Aissue+label%3Abug+label%3A%22topic%3A+torch_lib%22
• [torchlib] Fix scaled_dot_product_attention #968
• [torchlib] Fix hstack and vstack, atleast_*_sequence #960 @justinchuby
• [torchlib] Fix arange for int32 #974 @fatcat-z
• Support aten::div.Tensor_mode #980 Implement aten::div.Tensor_mode | feat(torchlib) #988

[justinchuby]

• repeat_interleave Is there a recommended implementation of yuv2RGB for the current torch? pytorch/pytorch#98077
  [ONNX] Simplify repeat_intereleave export for scalar-valued 'repeat' pytorch/pytorch#100575
• Implement _fft_* ops | feat(torchlib) #926 @fatcat-z
• Embedding_bag mismatch #1056 @xiaowuhu
• Batch norm: [torchlib] aten batch_norm ops have different size results on "CPU" vs "CUDA" #817 @titaiwangms
• [torchlib] Implement different padding modes #1138
• [torchlib] aten::index on boolean indicies #1141 @fatcat-z
• [torchlib] native_batch_norm needs more aten ops support #1140
• [torchlib] atleast_*d_sequence test fails #1007
• [torchlib] Implement various upsample functions #1159
• [torchlib] Implement aten::_scaled_dot_product_efficient_attention.default #1160
• Implement aten.var.correction #1174
• Implement aten.roll for complex inputs #1175
• aten_index_put doesn't support 1-D tensor #1212
• InstanceNorm always has training attribute set to True #1262 [torchlib] aten.instance_norm #1280
• torch.onnx.dynamo_export: Unsupported FX nodes: {'call_function': ['aten.pixel_shuffle.default']} pytorch/pytorch#120078
• [torchlib] torch.ops.aten.im2col #1290

[titaiwangms]

I found that PyTorch has been using this list: Core ATen ops:
https://github.com/pytorch/pytorch/blob/6eac3f45c7d6db45fa26507d0d8949ce6fb12dd7/torch/_decomp/__init__.py#L253

[justinchuby]

This is a list to be decomposed right?

[titaiwangms]

This is a list to be decomposed right?

Yes, some of these are not implemented and not decomposed, so users are getting unsupported errors. We could either decompose ATen Core, or implement them. pytorch/pytorch#125894.

[loretoparisi]

@justinchuby what is the aten::_native_multi_head_attention (from torch doc, not supported yet) support plan in the current opset? It seems in the opset-19 it is not supported yet I can see an error:

raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::_native_multi_head_attention' to ONNX opset version 17 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

[justinchuby]

@justinchuby what is the aten::_native_multi_head_attention (from torch doc, not supported yet) support plan in the current opset? It seems in the opset-19 it is not supported yet I can see an error:

raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::_native_multi_head_attention' to ONNX opset version 17 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: [https://github.com/pytorch/pytorch/issues.](https://github.com/pytorch/pytorch/issues.?q=sort%3Aupdated-desc+is%3Aissue+is%3Aopen)

Please leverage the torch.onnx.export(..., dynamo=True) option for this. It is available in the latest pytorch nightly build. Please let us know if you see any issues.