Repository: facebookresearch/NeuralCompression (language: python)

Issue:
lower_tail op
Approximates lower tail quantile for range coding.

```Python
lower_tail(distribution: torch.distributions.Distribution, tail_mass: float) -> float
```

For range coding of random variables, the distribution tails need special handling, because range coding can only handle alphabets with a finite number of symbols. This method returns a cut-off location for the lower tail, such that approximately `tail_mass` probability mass is contained in the tails (together). The tails are then handled by using the `overflow` functionality of the range coder implementation (using a Golomb-like universal code).

Prerequisites: 

- [x] #76 


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index bda7f303..2cf7745a 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -16,4 +16,5 @@ neuralcompression.functional
 .. autofunction:: log_ndtr
 .. autofunction:: log_survival_function
 .. autofunction:: log_cdf
+.. autofunction:: lower_tail
 .. autofunction:: lower_bound
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 6b885ef1..75b9073d 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -10,6 +10,7 @@
 from ._log_ndtr import log_ndtr
 from ._log_survival_function import log_survival_function
 from ._lower_bound import lower_bound
+from ._lower_tail import lower_tail
 from ._ndtr import ndtr
 from ._quantization_offset import quantization_offset
 from ._soft_round import (
diff --git a/neuralcompression/functional/_lower_tail.py b/neuralcompression/functional/_lower_tail.py
new file mode 100644
index 00000000..92cf39ac
--- /dev/null
+++ b/neuralcompression/functional/_lower_tail.py
@@ -0,0 +1,53 @@
+"""
+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 torch.distributions import Distribution
+
+from ._estimate_tails import estimate_tails
+from ._log_cdf import log_cdf
+
+
+def lower_tail(distribution: Distribution, tail_mass: float) -> Tensor:
+    """Approximates lower tail quantile for range coding.
+
+    For range coding of random variables, the distribution tails need special
+    handling, because range coding can only handle alphabets with a finite
+    number of symbols. This method returns a cut-off location for the lower
+    tail, such that approximately `tail_mass` probability mass is contained in
+    the tails (together). The tails are then handled by using the `overflow`
+    functionality of the range coder implementation (using a Golomb-like
+    universal code).
+
+    Args:
+        distribution: an object representing a continuous-valued random
+            variable.
+        tail_mass: desired probability mass for the tails.
+
+    Returns:
+        the approximate lower tail quantiles for each scalar distribution.
+    """
+    try:
+        x = distribution.icdf(torch.tensor([tail_mass / 2]))
+    except (AttributeError, NotImplementedError):
+        try:
+            x = estimate_tails(
+                lambda _x: log_cdf(_x, distribution),
+                math.log(tail_mass / 2),
+                distribution.batch_shape,
+            )
+        except NotImplementedError:
+            message = """
+            `distribution` must implement either `cdf` or `icdf`
+            """
+
+            raise NotImplementedError(message)
+
+    return x.detach()


Test patch:
diff --git a/tests/functional/test_lower_tail.py b/tests/functional/test_lower_tail.py
new file mode 100644
index 00000000..e7461465
--- /dev/null
+++ b/tests/functional/test_lower_tail.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 torch
+import torch.testing
+from torch.distributions import Normal, Uniform
+
+from neuralcompression.functional import lower_tail
+
+
+def test_lower_tail():
+    torch.testing.assert_close(
+        lower_tail(
+            Normal(0.0, 1.0),
+            0.5,
+        ),
+        torch.tensor([-0.6745]),
+    )
+
+    torch.testing.assert_close(
+        lower_tail(
+            Uniform(0.0, 1.0),
+            0.5,
+        ),
+        torch.tensor([0.25]),
+    )