xla_native_functions.yaml

backend: XLA
cpp_namespace: torch_xla
full_codegen:
  - acos
  - acosh
  - abs
  - all
  - all.dim
  - addcdiv
  - addcmul
  - amax
  - amin
  - any
  - any.dim
  - asin
  - asinh
  - atan
  - atanh
  - binary_cross_entropy
  - binary_cross_entropy_backward
  - bitwise_and.Tensor
  - bitwise_not
  - bitwise_or.Tensor
  - bitwise_xor.Tensor
  - ceil
  - cholesky
  - clamp.Tensor
  - clamp_max.Tensor
  - clamp_min.Tensor
  - cos
  - cosh
  - elu
  - eq.Scalar
  - eq.Tensor
  - erf
  - erfc
  - erfinv
  - exp
  - expm1
  - floor
  - frac
  - ge.Scalar
  - ge.Tensor
  - gt.Scalar
  - gt.Tensor
  - hardshrink
  - hardsigmoid
  - hardsigmoid_backward
  - hardswish
  - hardswish_backward
  - inverse
  - isnan
  - le.Scalar
  - le.Tensor
  - logdet
  - logical_and
  - logical_not
  - logical_or
  - logical_xor
  - log_sigmoid_backward
  - log_sigmoid_forward
  - lt.Scalar
  - lt.Tensor
  - maximum
  - minimum
  - ne.Scalar
  - ne.Tensor
  - reciprocal
  - relu
  - round
  - rsqrt
  - selu
  - sgn
  - sign
  - silu
  - silu_backward
  - sin
  - sinh
  - take
  - tan
  - tanh
  - tril
  - triu
  - trunc
ir_gen:
  - _adaptive_avg_pool2d
  - _adaptive_avg_pool2d_backward
  - _adaptive_avg_pool3d
  - _adaptive_avg_pool3d_backward
supported:
  - __ilshift__.Scalar
  - __ilshift__.Tensor
  - __irshift__.Scalar
  - __irshift__.Tensor
  - __lshift__.Scalar
  - __lshift__.Tensor
  - __rshift__.Scalar
  - __rshift__.Tensor
  - _adaptive_avg_pool2d
  - _adaptive_avg_pool2d_backward
  - _adaptive_avg_pool3d
  - _adaptive_avg_pool3d_backward
  - _amp_foreach_non_finite_check_and_unscale_
  - _amp_update_scale_
  - _copy_from
  - _copy_from_and_resize
  - _index_put_impl_
  - _local_scalar_dense
  - _log_softmax
  - _log_softmax_backward_data
  - _pack_padded_sequence
  - _softmax
  - _softmax_backward_data
  - _to_cpu
  - _trilinear
  - _unsafe_view
  - adaptive_max_pool2d
  - adaptive_max_pool2d_backward
  - add.Scalar
  - add.Tensor
  - addmm
  - alias
  - arange.start_out
  - argmax
  - argmin
  - as_strided
  - as_strided_
  - atan2
  - avg_pool2d
  - avg_pool2d_backward
  - avg_pool3d
  - avg_pool3d_backward
  - baddbmm
  - bernoulli
  - bernoulli_.float
  - bernoulli_.Tensor
  - binary_cross_entropy_with_logits
  - bmm
  - cat
  - celu
  - celu_
  - clamp
  - clamp_max
  - clamp_min
  - clone
  - constant_pad_nd
  - convolution_backward_overrideable
  - convolution_overrideable
  - cross
  - cumprod
  - cumsum
  - diag
  - diagonal
  - div.Scalar
  - div.Tensor
  - div.Tensor_mode
  - dot
  - elu_backward
  - embedding
  - embedding_dense_backward
  - empty.memory_format
  - empty_strided
  - expand
  - exponential_
  - eye.m_out
  - eye.out
  - fill_.Scalar
  - fill_.Tensor
  - flip
  - floor_divide
  - fmod.Scalar
  - fmod.Tensor
  - gather
  - gelu
  - gelu_backward
  - hardshrink_backward
  - hardtanh
  - hardtanh_backward
  - index.Tensor
  - index_add
  - index_copy
  - index_fill_.int_Scalar
  - index_fill_.int_Tensor
  - index_put_
  - index_select
  - kl_div
  - kthvalue
  - leaky_relu
  - leaky_relu_backward
  - lerp.Scalar
  - lerp.Tensor
  - linspace
  - log
  - log1p
  - log2
  - log10
  - logsumexp
  - masked_fill_.Scalar
  - masked_fill_.Tensor
  - masked_scatter_
  - masked_select
  - max
  - max.dim
  - max.dim_max
  - max_pool2d_with_indices
  - max_pool2d_with_indices_backward
  - max_pool3d_with_indices
  - max_pool3d_with_indices_backward
  - max_unpool2d
  - max_unpool3d
  - mean
  - mean.dim
  - min
  - min.dim
  - min.dim_min
  - mish
  - mm
  - mse_loss
  - mse_loss_backward
  - mul.Scalar
  - mul.Tensor
  - mv
  - mv.out
  - nan_to_num
  - native_batch_norm
  - native_batch_norm_backward
  - neg
  - nll_loss2d_backward
  - nll_loss2d_forward
  - nll_loss_backward
  - nll_loss_forward
  - nonzero
  - norm.Scalar
  - norm.ScalarOpt_dim
  - norm.ScalarOpt_dim_dtype
  - norm.ScalarOpt_dtype
  - normal.float_Tensor
  - normal.Tensor_float
  - normal.Tensor_Tensor
  - normal_
  - permute
  - pow.Scalar
  - pow.Tensor_Scalar
  - pow.Tensor_Tensor
  - prelu
  - prod
  - prod.dim_int
  - put_
  - qr
  - random_
  - random_.from
  - random_.to
  - reflection_pad2d
  - reflection_pad2d_backward
  - remainder.Scalar
  - remainder.Tensor
  - repeat
  - replication_pad1d
  - replication_pad1d_backward
  - replication_pad2d
  - replication_pad2d_backward
  - resize_
  - roll
  - rrelu_with_noise
  - rrelu_with_noise_backward
  - rsub.Scalar
  - rsub.Tensor
  - scatter.reduce
  - scatter.src
  - scatter.value
  - scatter.value_reduce
  - scatter_add
  - select.int
  - selu_
  - sigmoid
  - sigmoid_backward
  - slice.Tensor
  - slogdet
  - smooth_l1_loss
  - smooth_l1_loss_backward
  - softplus
  - softplus_backward
  - softshrink
  - softshrink_backward
  - sort
  - sort.stable
  - split.Tensor
  - split_with_sizes
  - sqrt
  - squeeze
  - squeeze.dim
  - squeeze_
  - squeeze_.dim
  - stack
  - std
  - std.correction
  - std.dim
  - std_mean.correction
  - sub.Scalar
  - sub.Tensor
  - sum
  - sum.dim_IntList
  - svd
  - symeig
  - t
  - t_
  - tanh_backward
  - threshold
  - threshold_backward
  - topk
  - trace
  - transpose.int
  - transpose_
  - triangular_solve
  - unbind.int
  - uniform_
  - unsqueeze
  - unsqueeze_
  - upsample_bilinear2d
  - upsample_bilinear2d_backward
  - upsample_nearest2d
  - upsample_nearest2d.vec
  - upsample_nearest2d_backward
  - upsample_nearest2d_backward.vec
  - var.correction
  - var_mean.correction
  - view
  - where.self
  - xlogy.Tensor
  - zero_
  # Below are all operators that are "composite" in core,
  # but require us to explicitly re-enable functionalization in order to use them.
  # Why? These operators are all CompositeExplicitAutograd, which mean that they run
  # after functionalization,
  # but their implementations call view operators (which we need to functionalize away).
  - block_diag
  - slice_backward
  - diagonal_backward
  - new_empty_strided
  - narrow_copy
  - pixel_shuffle
  - pixel_unshuffle
  - select_backward
  - linalg_pinv.atol_rtol_tensor
  # The same applies to these ops, but we already have direct lowerings for them
  # - _trilinear
  # - logsumexp.out
symint:
  - empty.memory_format
  - empty_strided
  - expand
  - new_empty_strided
  - view
autograd:
  - einsum
  - max_pool2d
  - max_pool3d
  - native_layer_norm
  - native_group_norm