Repository: pytorch
Issue: Rollup: forward-mode AD operator coverage
Labels: module: autograd,triaged,actionable,module: forward ad

NOTE: See also the [tracker](https://github.com/pytorch/pytorch/issues/75432) for forward-over-reverse formulas:  Sometimes implementing a forward AD formula will also implicitly support its forward-over-reverse formula, make sure to comment/mark as completed on that issue as well to keep it up to date.

These are ops that currently have `supports_forward_ad=False` (usually implicitly) in their OpInfo. Because not all ops have OpInfos yet, some ops that don't support forward AD may be missing from this list. **Please file an issue or comment here if there is an op you'd like to see supported, but is not on this list**. A complete list of formulas is included at the bottom.

For a short tutorial on adding a forward AD formula see: https://github.com/pytorch/pytorch/issues/64545

If you take a brief look at the list below, you'll notice that formulas in `derivatives.yaml` don't generally correspond one-to-one with OpInfo entries, so it is sometimes hard to know just from staring at the op name, which formula is needed to support it. To figure out what formula corresponds to a particular op, set `supports_forward_ad=True` in the OpInfo and run the test `test_forward_mode_AD`, the error will tell you which formula is needed, e.g. `forward ad is not implemented for <formula name>`. Sometimes a single formula will support multiple ops. After implementing the formula in `dervatives.yaml`, it may be useful to run all the forward AD tests to reveal if any more ops have been supported by side effect.

High priority (*in order of priority*; please start working from the top) 
- [x] nn.functional.binary_cross_entropy (DO THIS FIRST; needs OpInfo) https://github.com/pytorch/pytorch/pull/77755
- [x] nn.functional.binary_cross_entropy_with_logits (DO THIS FIRST; needs OpInfo) (already there, but https://github.com/pytorch/pytorch/pull/77755 simplifies implementation)
- [x] clamp https://github.com/pytorch/pytorch/pull/74042
- [x] log_softmax https://github.com/pytorch/pytorch/pull/73741
- [x] norm https://github.com/pytorch/pytorch/pull/70253, https://github.com/pytorch/pytorch/pull/74205
- [x] dist https://github.com/pytorch/pytorch/pull/74205
- [x] atan2 https://github.com/pytorch/pytorch/pull/75027
- [x] nn.functional.embedding @Lezcano 
- [x] nn.functional.glu https://github.com/pytorch/pytorch/pull/77186
- [x] nn.functional.prelu @Lezcano 
- [x] nn.functional.multi_head_attention_forward
- [x] nn.functional.dropout (@zou3519)
- [x] nn.functional.{dropout2d, dropout3d}
- [x] nn.functional.l1_loss
- [x] nn.functional.smooth_l1_loss
- [ ] nn.functional.conv3d (needs OpInfo)
- [x] fmod https://github.com/pytorch/pytorch/pull/69908
- [x] softmax https://github.com/pytorch/pytorch/pull/73741
- [x] nn.functional.softmin https://github.com/pytorch/pytorch/pull/73741
- [x] nn.functional.cross_entropy https://github.com/pytorch/pytorch/pull/73741
- [x] nn.functional.normalize https://github.com/pytorch/pytorch/pull/74205
- [x] nn.functional.bilinear @Lezcano 
- [x] nn.functional.rrelu @Lezcano 
- [x] nn.functional.logsigmoid @Lezcano 
- [x] `__rsub__` https://github.com/pytorch/pytorch/pull/75326
- [x] rsub https://github.com/pytorch/pytorch/pull/75326
- [x] polar https://github.com/pytorch/pytorch/pull/75326
- [x] `__getitem__` https://github.com/pytorch/pytorch/pull/69909
- [x] put
- [ ] cdist
- [x] renorm https://github.com/pytorch/pytorch/pull/100798
- [x] logsumexp https://github.com/pytorch/pytorch/pull/73741
- [x] logcumsumexp https://github.com/pytorch/pytorch/pull/100629
- [ ] nn.functional.embedding_bag
- [ ] nn.functional.grid_sample
- [x] amax https://github.com/pytorch/pytorch/pull/77975
- [x] amin https://github.com/pytorch/pytorch/pull/77975
- [x] nanmean https://github.com/pytorch/pytorch/pull/77975
- [x] nansum https://github.com/pytorch/pytorch/pull/77975
- [ ] nn.functional.ctc_loss (?)
- [x] nn.functional.nll_loss https://github.com/pytorch/pytorch/pull/73741
- [x] nn.functional.pairwise_distance https://github.com/pytorch/pytorch/pull/74205
- [x] nn.functional.mse_loss https://github.com/pytorch/pytorch/pull/71026
- [x] nn.functional.huber_loss https://github.com/pytorch/pytorch/pull/71026
- [x] nn.functional.kl_div https://github.com/pytorch/pytorch/pull/71026

Linalg ops (old tracker: https://github.com/pytorch/pytorch/issues/64545)
- [x] linalg.det https://github.com/pytorch/pytorch/pull/79487
- [ ] linalg.cond
- [x] linalg.eigvals https://github.com/pytorch/pytorch/pull/70527
- [x] linalg.lstsq (already done, see old tracker - need to filter out non-grad oriented OpInfo)
- [x] linalg.norm https://github.com/pytorch/pytorch/pull/74205
- [x] linalg.matrix_norm (TODO: @soulitzer, investigate failure)
- [x] linalg.slogdet https://github.com/pytorch/pytorch/pull/79743
- [x] linalg.vector_norm https://github.com/pytorch/pytorch/pull/74205
- [x] cholesky_inverse https://github.com/pytorch/pytorch/pull/75033
- [x] lu_unpack https://github.com/pytorch/pytorch/pull/67833
- [x] svd  https://github.com/pytorch/pytorch/pull/70253
- [x] linalg.svd  https://github.com/pytorch/pytorch/pull/70253
- [x] linalg.svdvals https://github.com/pytorch/pytorch/pull/70253
- [x] svd_lowrank (composite op, supported by linalg.svd)
- [x] pca_lowrank (composite op, supported by linalg.svd)
- [x] logdet https://github.com/pytorch/pytorch/pull/79743

Low frequency:
- [ ] igamma
- [ ] igammac
- [ ] special.zeta
- [ ] to_sparse
- [ ] nn.functional.one_hot

FFT ops (these correspond to 3 formulas in `derivatives.yaml`: _fft_r2c, _fft_c2r, _fft_c2c): 
PR in progress: https://github.com/pytorch/pytorch/pull/75326
- [x] fft.fft
- [x] fft.fft2
- [x] fft.fftn
- [x] fft.hfft
- [x] fft.hfft2
- [x] fft.hfftn
- [x] fft.rfft
- [x] fft.rfft2
- [x] fft.rfftn
- [x] fft.ifft
- [x] fft.ifft2
- [x] fft.ifftn
- [x] fft.ihfft
- [x] fft.ihfft2
- [x] fft.ihfftn
- [x] fft.irfft
- [x] fft.irfft2
- [x] fft.irfftn
- [x] stft
- [x] istft

Masked ops (these don't have their own formulas, just implement the formula for the non-masked version and you get covered here for free):
- [x] _masked.amax https://github.com/pytorch/pytorch/pull/77975
- [x] _masked.amin https://github.com/pytorch/pytorch/pull/77975
- [x] _masked.norm https://github.com/pytorch/pytorch/pull/74205
- [x] _masked.softmax https://github.com/pytorch/pytorch/pull/73741
- [x] _masked.log_softmax https://github.com/pytorch/pytorch/pull/73741
- [x] _masked.softmin https://github.com/pytorch/pytorch/pull/73741
- [x] _masked.normalize https://github.com/pytorch/pytorch/pull/74205

Deprecated ops (https://github.com/pytorch/pytorch/issues/55953):
- cholesky
- solve
- symeig
- eig


It may not be entirely useful, but expand the below list (TODO: not up to date) to see a list of remaining formulas. Implementing a formula that ends in  `_backward`  will support forward-over-reverse autograd. 

<details>
<summary>
[click to expand]
</summary>

```
_sparse_addmm
affine_grid_generator
acosh_
asinh_
atanh_
atan2
cholesky
cholesky_inverse
clamp.Tensor
_coalesce
polar
logcumsumexp
_ctc_loss
_det_lu_based_helper
dist
_fused_dropout
native_dropout
native_dropout_backward
eig
fake_quantize_per_tensor_affine_cachemask
_fake_quantize_per_tensor_affine_cachemask_tensor_qparams
_fake_quantize_learnable_per_tensor_affine
fake_quantize_per_channel_affine_cachemask
_fake_quantize_learnable_per_channel_affine
_fused_moving_avg_obs_fq_helper
geqrf
grid_sampler_2d
grid_sampler_3d
_grid_sampler_2d_cpu_fallback
igamma
igammac
index.Tensor https://github.com/pytorch/pytorch/pull/69909
special_zeta
special_zeta.self_scalar
special_zeta.other_scalar
logdet
logsumexp
lstsq
lu_unpack
amax
amin
native_batch_norm_backward
native_layer_norm_backward
nextafter
norm.Scalar
norm.ScalarOpt_dim
norm.ScalarOpt_dtype
norm.ScalarOpt_dim_dtype
linalg_vector_norm
_pdist_forward
_pdist_backward
_euclidean_dist
_cdist_forward
_cdist_backward
ormqr
put_
renorm
slogdet
linalg_slogdet
solve
std_mean.correction
rsub.Tensor
rsub.Scalar
nansum
nansum.dim_IntList
_svd_helper
symeig
one_hot
to_dense
to_sparse
to_sparse.sparse_dim
to_mkldnn
_unique
unique_dim
unique_consecutive
unique_dim_consecutive
_unique2
var_mean.correction
_weight_norm_cuda_interface
sparse_mask
_sparse_coo_tensor_with_dims_and_tensors
_sparse_sum.dim
_standard_gamma
_standard_gamma_grad
values
_trilinear
binary_cross_entropy
binary_cross_entropy_backward
embedding
embedding_dense_backward
_embedding_bag
_embedding_bag_dense_backward
embedding_renorm_
multilabel_margin_loss_forward
nll_loss_forward
nll_loss2d_forward
elu_
celu_
glu
log_sigmoid_forward
_log_softmax
_sparse_log_softmax
prelu
prelu_backward
rrelu_with_noise
rrelu_with_noise_
_softmax
_sparse_softmax
_sparse_sparse_matmul
_upsample_bicubic2d_aa.vec
elu_backward
glu_backward
kl_div_backward
l1_loss_backward
log_sigmoid_backward
_log_softmax_backward_data
leaky_relu_backward
max_unpool2d_backward
mse_loss_backward
nll_loss_backward
nll_loss2d_backward
rrelu_with_noise_backward
smooth_l1_loss_backward
huber_loss_backward
softplus_backward
_softmax_backward_data
soft_margin_loss_backward
_upsample_bicubic2d_aa_backward.vec
_cudnn_ctc_loss
cudnn_grid_sampler
cudnn_affine_grid_generator
cudnn_batch_norm_backward
_cudnn_rnn
_cudnn_rnn_backward
miopen_batch_norm_backward
miopen_rnn
miopen_rnn_backward
mkldnn_linear
mkldnn_max_pool2d
mkldnn_max_pool3d
mkldnn_adaptive_avg_pool2d
_mkldnn_reshape
_fft_r2c
_fft_c2r
_fft_c2c
_thnn_fused_lstm_cell
_thnn_fused_gru_cell
_pack_padded_sequence
segment_reduce
_pin_memory
_test_warn_in_autograd
```

</details>

A follow up tracker can be created to track the completion of any remaining formulas if necessary. Alternatively, any additions to the list above are also welcome.

See also: Tracker for `supports_fwgrad_bwgrad=False` for forward-over-reverse Hessian coverage (TODO).

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7