Repository: facebookresearch/NeuralCompression (language: python)

Issue:
log_expm1 op
Logarithm of :math:`e^{x} - 1`.


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index 1e275d53..b41f241f 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -19,3 +19,4 @@ neuralcompression.functional
 .. autofunction:: lower_tail
 .. autofunction:: lower_bound
 .. autofunction:: upper_bound
+.. autofunction:: log_expm1
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 33fee15d..3c8d8249 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -7,6 +7,7 @@
 
 from ._estimate_tails import estimate_tails
 from ._log_cdf import log_cdf
+from ._log_expm1 import log_expm1
 from ._log_ndtr import log_ndtr
 from ._log_survival_function import log_survival_function
 from ._lower_bound import lower_bound
diff --git a/neuralcompression/functional/_log_expm1.py b/neuralcompression/functional/_log_expm1.py
new file mode 100644
index 00000000..164122b5
--- /dev/null
+++ b/neuralcompression/functional/_log_expm1.py
@@ -0,0 +1,34 @@
+"""
+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
+
+
+def log_expm1(
+    x: Tensor,
+) -> Tensor:
+    r"""Logarithm of :math:`e^{x} - 1`.
+
+    If ``x`` is large, ``torch.exp(x)`` returns ``float("inf")`` whereas:
+
+    .. math::
+        \log e^{x} - 1 \approx x
+
+    Therefore, an approximation for :math:`x > 15` is used, such that the
+    output is not infinity for all positive values of ``x``.
+
+    Args:
+        x: the input tensor.
+
+    Returns:
+        the logarithm of :math:`e^{x} - 1`.
+    """
+    return torch.where(
+        x < 15.0,
+        torch.log(torch.expm1(torch.minimum(x, torch.tensor(15.0)))),
+        x,
+    )


Test patch:
diff --git a/tests/functional/test_log_expm1.py b/tests/functional/test_log_expm1.py
new file mode 100644
index 00000000..4c384821
--- /dev/null
+++ b/tests/functional/test_log_expm1.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 torch
+import torch.testing
+
+from neuralcompression.functional import log_expm1
+
+
+def test_log_expm1():
+    rng = numpy.random.default_rng(0xFEEEFEEE)
+
+    x = rng.random((32,), dtype=numpy.float)
+
+    actual = log_expm1(torch.tensor(x, dtype=torch.float))
+
+    assert torch.isfinite(actual).all()
+
+    torch.testing.assert_close(
+        actual,
+        torch.tensor(numpy.log(numpy.expm1(x)), dtype=torch.float),
+    )