Repository: facebookresearch/NeuralCompression (language: python)

Issue:
survival_function op
Survival function of `x`. Generally defined as `1 - distribution.cdf(x)`.

```Python
survival_function(x: torch.Tensor, distribution: torch.distributions.Distribution) -> torch.Tensor
```

Prerequisites 

- [x] #78 


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index d1a3383e..9ff0b9b4 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -12,4 +12,5 @@ neuralcompression.functional
 .. autofunction:: dense_image_warp
 .. autofunction:: quantization_offset
 .. autofunction:: ndtr
+.. autofunction:: survival_function
 .. autofunction:: log_ndtr
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 70f87751..0a5d5246 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -14,6 +14,7 @@
     soft_round_conditional_mean,
     soft_round_inverse,
 )
+from ._survival_function import survival_function
 from .complexity import count_flops
 from .distortion import (
     learned_perceptual_image_patch_similarity,
diff --git a/neuralcompression/functional/_survival_function.py b/neuralcompression/functional/_survival_function.py
new file mode 100644
index 00000000..0caa6bcb
--- /dev/null
+++ b/neuralcompression/functional/_survival_function.py
@@ -0,0 +1,35 @@
+"""
+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 torch
+from torch import Tensor
+from torch.distributions import Distribution, Normal
+
+from ._ndtr import ndtr
+
+
+def survival_function(
+    x: Tensor,
+    distribution: Distribution,
+) -> Tensor:
+    """Survival function of ``x``. Generally defined as
+    ``1 - distribution.cdf(x)``.
+
+    Args:
+        x: the input tensor.
+        distribution: an object representing a continuous-valued random
+        variable.
+
+    Returns:
+        the survival function of ``x`` for the continuous-valued random
+        variable ``distribution``.
+    """
+    if isinstance(distribution, Normal):
+        standardized = torch.div((x - distribution.loc), distribution.scale)
+
+        return ndtr(-standardized)
+
+    return 1.0 - distribution.cdf(x)


Test patch:
diff --git a/tests/functional/test_survival_function.py b/tests/functional/test_survival_function.py
new file mode 100644
index 00000000..a8959579
--- /dev/null
+++ b/tests/functional/test_survival_function.py
@@ -0,0 +1,30 @@
+"""
+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
+import scipy.stats
+import torch
+import torch.testing
+from torch.distributions import Normal
+
+from neuralcompression.functional import survival_function
+
+
+def test_survival_function():
+    random = numpy.random.default_rng(0xDEADBEEF)
+
+    batch_size = 32
+
+    loc = torch.tensor([random.standard_normal(batch_size)])
+
+    scale = torch.tensor([random.random(batch_size) + 1.0])
+
+    x = torch.linspace(-8.0, 8.0, batch_size).to(torch.float64)
+
+    torch.testing.assert_allclose(
+        survival_function(x, Normal(loc, scale)),
+        scipy.stats.norm(loc, scale).sf(x),
+    )