Repository: facebookresearch/NeuralCompression (language: python)

Issue:
log_survival_function op
Logarithm of x for a distribution’s survival function.

```Python
log_survival_function(x: torch.Tensor, distribution: torch.distributions.Distribution) -> torch.Tensor
```

Prerequisites

- [x] #79 


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index 3a91d01f..dbcb0569 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -14,4 +14,5 @@ neuralcompression.functional
 .. autofunction:: ndtr
 .. autofunction:: survival_function
 .. autofunction:: log_ndtr
+.. autofunction:: log_survival_function
 .. autofunction:: log_cdf
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 6b402c95..4d5457dd 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -8,6 +8,7 @@
 from ._estimate_tails import estimate_tails
 from ._log_cdf import log_cdf
 from ._log_ndtr import log_ndtr
+from ._log_survival_function import log_survival_function
 from ._ndtr import ndtr
 from ._quantization_offset import quantization_offset
 from ._soft_round import (
diff --git a/neuralcompression/functional/_log_survival_function.py b/neuralcompression/functional/_log_survival_function.py
new file mode 100644
index 00000000..9d5bcf2d
--- /dev/null
+++ b/neuralcompression/functional/_log_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 ._log_ndtr import log_ndtr
+
+
+def log_survival_function(
+    x: Tensor,
+    distribution: Distribution,
+) -> Tensor:
+    """Logarithm of ``x`` for a distribution’s survival function.
+
+    Args:
+        x: the input tensor.
+        distribution: an object representing a continuous-valued random
+        variable.
+
+    Returns:
+        the logarithm of ``x`` for a distribution’s survival function.
+    """
+
+    if isinstance(distribution, Normal):
+        standardized = torch.div((x - distribution.loc), distribution.scale)
+
+        return log_ndtr(-standardized)
+
+    return torch.log1p(distribution.cdf(-x))


Test patch:
diff --git a/tests/functional/test_log_survival_function.py b/tests/functional/test_log_survival_function.py
new file mode 100644
index 00000000..1df19309
--- /dev/null
+++ b/tests/functional/test_log_survival_function.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 numpy
+import scipy.stats
+import torch
+import torch.testing
+from torch.distributions import Normal
+
+from neuralcompression.functional import log_survival_function
+
+
+def test_log_survival_function():
+    rng = numpy.random.default_rng(0xFEEEFEEE)
+
+    batch_size = 32
+
+    loc = rng.standard_normal(batch_size, dtype=numpy.float)
+
+    scale = rng.random(batch_size, dtype=numpy.float) + 1.0
+
+    x = numpy.linspace(-8.0, 8.0, batch_size, dtype=numpy.float)
+
+    torch.testing.assert_allclose(
+        log_survival_function(
+            torch.tensor(x, dtype=torch.float),
+            Normal(
+                torch.tensor(loc, dtype=torch.float),
+                torch.tensor(scale, dtype=torch.float),
+            ),
+        ),
+        scipy.stats.norm(loc, scale).logsf(x),
+    )