Repository: facebookresearch/NeuralCompression (language: python)

Issue:
log_ndtr op
Logarithm of the normal cumulative distribution function (CDF).

```Python
log_ndtr(x: torch.Tensor) -> torch.Tensor
```

Prerequisites 

- [x] #78 


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index 6b2a9f2e..e5a7c5b3 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -11,3 +11,4 @@ neuralcompression.functional
 .. autofunction:: optical_flow_to_color
 .. autofunction:: dense_image_warp
 .. autofunction:: ndtr
+.. autofunction:: log_ndtr
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 231f0e4c..f01237ca 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -6,6 +6,7 @@
 """
 
 from ._estimate_tails import estimate_tails
+from ._log_ndtr import log_ndtr
 from ._ndtr import ndtr
 from ._soft_round import (
     soft_round,
diff --git a/neuralcompression/functional/_log_ndtr.py b/neuralcompression/functional/_log_ndtr.py
new file mode 100644
index 00000000..7d046e7d
--- /dev/null
+++ b/neuralcompression/functional/_log_ndtr.py
@@ -0,0 +1,59 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import math
+
+import torch
+from torch import Tensor
+
+from ._ndtr import ndtr
+
+
+def _log_ndtr_asymptotic_series(x: Tensor, series_order: int = 3) -> Tensor:
+    t1 = -0.5 * (math.log(2 * math.pi) + x ** 2) - torch.log(-x)
+    t2 = torch.zeros_like(x)
+    value_even_power = (x ** 2).clone()
+    double_fac = 1
+    multiplier = -1
+
+    for n in range(1, series_order + 1):
+        t2.add_(multiplier * double_fac / value_even_power)
+        value_even_power.mul_(x ** 2)
+        double_fac *= 2 * n - 1
+        multiplier *= -1
+
+    return t1 + torch.log1p(t2)
+
+
+def log_ndtr(x: Tensor) -> Tensor:
+    """Logarithm of the normal cumulative distribution function (CDF).
+
+    Args:
+        x: the input tensor.
+
+    Returns:
+        the log of the area under the standard Normal probability density
+            function (PDF), integrated from minus infinity to ``x``.
+    """
+    if x.dtype == torch.float32:
+        m, n = -10.0, 5.0
+    elif x.dtype == torch.float64:
+        m, n = -20.0, 8.0
+    else:
+        raise TypeError(
+            f"`{log_ndtr.__name__}` doesn’t support {x.dtype}",
+        )
+
+    return torch.where(
+        x > n,
+        torch.log1p(-ndtr(-x)),
+        torch.where(
+            x >= m,
+            torch.log(ndtr(x)),
+            _log_ndtr_asymptotic_series(x),
+        ),
+    )


Test patch:
diff --git a/tests/functional/test_log_ndtr.py b/tests/functional/test_log_ndtr.py
new file mode 100644
index 00000000..de7ac084
--- /dev/null
+++ b/tests/functional/test_log_ndtr.py
@@ -0,0 +1,27 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import numpy.random
+import scipy.special
+import torch
+import torch.testing
+
+from neuralcompression.functional import log_ndtr
+
+
+def test_log_ndtr():
+    rng = numpy.random.default_rng(0xDEADBEEF)
+
+    x = rng.random((32,), dtype=numpy.float)
+
+    actual = log_ndtr(torch.tensor(x, dtype=torch.float))
+
+    assert torch.isfinite(actual).all()
+
+    torch.testing.assert_close(
+        actual,
+        torch.tensor(scipy.special.log_ndtr(x), dtype=torch.float),
+    )