Hi Samuel!



 Samue1:

Does this also work for batched versions of S?


No.  If you had tried it, you would have discovered that torch.outer()
does not accept multidimensional tensors.


Is the result correct if I use
J = torch.diag_embed(S) - torch.outer(S, S)


No, this will throw an error (because you pass a multidimensional
tensor to torch.outer()).
You can, however, use pytorch’s swiss army knife of tensor
multiplication functions to construct a batch version of outer:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> S = torch.arange (6).reshape (2, 3).float()
>>> S
tensor([[0., 1., 2.],
        [3., 4., 5.]])
>>> torch.diag_embed (S) - torch.einsum ('ij, ik -> ijk', S, S)
tensor([[[  0.,   0.,   0.],
         [  0.,   0.,  -2.],
         [  0.,  -2.,  -2.]],

        [[ -6., -12., -15.],
         [-12., -12., -20.],
         [-15., -20., -20.]]])

(As an aside, none of this has anything to do with the title you gave
this thread, namely “Jacobian of Softmax.”)
Best.
K. Frank