Repository: pytorch
Issue: Add Lambert W function as torch.special.lambertw
Labels: triaged,module: special,topic: new features

### 🚀 The feature, motivation and pitch

Implement the Lambert W function as a core mathematical function available in pytorch.  Suggest to put this in `torch.special.lambertw` to follow pattern in [`scipy.special.lambertw`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.lambertw.html) and [`tfp.special.lambertw`](https://www.tensorflow.org/probability/api_docs/python/tfp/math/lambertw). 

Lambert W function ([wiki](https://en.wikipedia.org/wiki/Lambert_W_function), and [Corless et al, 1996](https://cs.uwaterloo.ca/research/tr/1993/03/W.pdf)) is the solution to the equation `x * exp(x) = z`, ie `W(z)=x` with multiple branches (real and complex).  It has many applications in physics & math and skewed/heavy-tailed distribution modeling; recently this also gained attention in the deep learning literature as a way to add Gaussianization (outlier dampeninig / remove heavy tails) layers ([Goerg 2011](https://projecteuclid.org/journals/annals-of-applied-statistics/volume-5/issue-3/Lambert-W-random-variablesa-new-family-of-generalized-skewed-distributions/10.1214/11-AOAS457.full) & [Goerg 2015](https://www.hindawi.com/journals/tswj/2015/909231/), [Wiese et al, 2019](https://www.tandfonline.com/doi/abs/10.1080/14697688.2020.1730426), [Li 2023](https://openreview.net/pdf?id=OXP9Ns0gnIq), amongst others). 

Adding it here as a core `pytorch.special.lambertw` functionality would open up ability for pytorch developers to implement any advances involving Lambert W functions; currently blocked by lack of implementation compared to other standard frameworks/libraries (Boost C++, TensorFlow, scipy, R, Matlab, Julia):

* Boost C++ implementation: https://www.boost.org/doc/libs/1_69_0/libs/math/doc/html/math_toolkit/lambert_w.html
* `sklearn` like `gaussianize` Transformer: https://github.com/gregversteeg/gaussianize
* `LambertW` R package: https://github.com/gmgeorg/LambertW; based on C++ implementation in `lamW` (https://github.com/cran/lamW/blob/master/src/lambertW.cpp)
* TensorFlow implementation of `lambertw`: https://www.tensorflow.org/probability/api_docs/python/tfp/math/lambertw
* MATLAB (https://www.mathworks.com/help/symbolic/lambertw.html)
* Julia (https://docs.juliahub.com/LambertW/7mpiq/0.4.5/autodocs/)

Particularly it would unblock implemented Lambert W x F distributions in `torch.distributions.*` (akin to their [TensorFlow Probability `LambertWDistribution` counterpart](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/LambertWDistribution)) and subsequently in `pyro.ai`, as well as Gaussianization layers.

**Fun fact**: for gradient based computations Lambert W function has the nice property that its derivative can be expressed in terms of the original function value (ie fast computation of derivatives): `W'(z) = W(z) / (z * (1 + W(z))`. See eq (3.2) in https://cs.uwaterloo.ca/research/tr/1993/03/W.pdf


### Alternatives

none; the Lambert W function is a core mathematical function that's missing in `pytorch`. 

### Additional context

Basic implementations of Lambert W already exist here:

https://github.com/thibsej/unbalanced-ot-functionals/blob/13f2203b3993d973f929578085ea458c5c1a7a78/common/torch_lambertw.py

https://github.com/AminJun/BreakingCertifiableDefenses/blob/cc469fa48f7efba21f3584e233c4db0c9a4856c1/RandomizedSmoothing/projected_sinkhorn/lambertw.py

EDIT: I took a stab at this and implemented a native `torch` version of `lambertw` translating https://www.tensorflow.org/probability/api_docs/python/tfp/math/lambertw with e2e testing (asserting equality to scipy implementation and known identities of W).  Just need some guidance on what constitutes an e2e fully implemented new `torch.special.lambertw` function (gradients decorators, ...); best a recent example PR that implemented a different basic core mathematical function then I can pattern match it.

cc @mruberry @kshitij12345