Repository: facebookresearch/NeuralCompression (language: python)

Issue:
ndtr op
Normal cumulative distribution function (CDF).

```Python
ndtr(x: torch.Tensor) -> torch.Tensor
```

Returns the area under the standard Normal probability density function (PDF), integrated from negative infinity to `x`.



Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index feda29d8..6b2a9f2e 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -10,3 +10,4 @@ neuralcompression.functional
 .. autofunction:: hsv2rgb
 .. autofunction:: optical_flow_to_color
 .. autofunction:: dense_image_warp
+.. autofunction:: ndtr
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index cc2316e1..231f0e4c 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -6,6 +6,7 @@
 """
 
 from ._estimate_tails import estimate_tails
+from ._ndtr import ndtr
 from ._soft_round import (
     soft_round,
     soft_round_conditional_mean,
diff --git a/neuralcompression/functional/_ndtr.py b/neuralcompression/functional/_ndtr.py
new file mode 100644
index 00000000..882d7560
--- /dev/null
+++ b/neuralcompression/functional/_ndtr.py
@@ -0,0 +1,36 @@
+"""
+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
+
+
+def ndtr(x: Tensor) -> Tensor:
+    """The normal cumulative distribution function (CDF).
+
+    Args:
+        x: the input tensor.
+
+    Returns:
+        the area under the standard Normal probability density function (PDF),
+            integrated from negative infinity to ``x``.
+    """
+    x = x.to(torch.float) * math.sqrt(0.5)
+
+    y = 0.5 * torch.erfc(abs(x))
+
+    return torch.where(
+        abs(x) < math.sqrt(0.5),
+        0.5 + 0.5 * torch.erf(x),
+        torch.where(
+            x > 0,
+            1 - y,
+            y,
+        ),
+    )


Test patch:
diff --git a/tests/functional/test_ndtr.py b/tests/functional/test_ndtr.py
new file mode 100644
index 00000000..79d79453
--- /dev/null
+++ b/tests/functional/test_ndtr.py
@@ -0,0 +1,29 @@
+"""
+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 ndtr
+
+
+def test_ndtr():
+    torch.testing.assert_equal(
+        ndtr(torch.tensor([0.0])),
+        torch.tensor([0.5]),
+    )
+
+    rng = numpy.random.default_rng(0xDEADBEEF)
+
+    x = rng.random((32,), dtype=numpy.float)
+
+    torch.testing.assert_close(
+        ndtr(torch.tensor(x, dtype=torch.float)),
+        torch.tensor(scipy.special.ndtr(x), dtype=torch.float),
+    )