make tsnex differentiable

[alonfnt]

Right now, tsnex is not differentiable as it relies on a while loop to compute the perplexity.. It does use autograd for minimising the loss nevertheless. This while loop can probably be fixed by defining the custom rules of the backwards pass using the implicit function theorem.

That would allow to apply grad on where there's the transform call inside.