# Copyright (c) Meta Platforms, Inc. and 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 neuralcompression.functional import soft_round, soft_round_inverse


def test_soft_round_inverse():
    x = torch.linspace(-2.0, 2.0, 50)

    torch.testing.assert_close(
        x,
        soft_round_inverse(x, alpha=1e-13),
    )

    x = torch.tensor([-1.25, -0.75, 0.75, 1.25])

    torch.testing.assert_close(
        x,
        soft_round_inverse(soft_round(x, alpha=2.0), alpha=2.0),
    )

    for offset in range(-5, 5):
        x = torch.linspace(offset + 0.001, offset + 0.999, 100)

        torch.testing.assert_close(
            torch.ceil(x) - 0.5,
            soft_round_inverse(x, alpha=5000.0),
            atol=0.001,
            rtol=0.002,
        )

[((307, 336), 'torch.linspace', 'torch.linspace', (['(-2.0)', '(2.0)', '(50)'], {}), '(-2.0, 2.0, 50)\n', (321, 336), False, 'import torch\n'), ((440, 480), 'torch.tensor', 'torch.tensor', (['[-1.25, -0.75, 0.75, 1.25]'], {}), '([-1.25, -0.75, 0.75, 1.25])\n', (452, 480), False, 'import torch\n'), ((389, 423), 'neuralcompression.functional.soft_round_inverse', 'soft_round_inverse', (['x'], {'alpha': '(1e-13)'}), '(x, alpha=1e-13)\n', (407, 423), False, 'from neuralcompression.functional import soft_round, soft_round_inverse\n'), ((641, 692), 'torch.linspace', 'torch.linspace', (['(offset + 0.001)', '(offset + 0.999)', '(100)'], {}), '(offset + 0.001, offset + 0.999, 100)\n', (655, 692), False, 'import torch\n'), ((552, 576), 'neuralcompression.functional.soft_round', 'soft_round', (['x'], {'alpha': '(2.0)'}), '(x, alpha=2.0)\n', (562, 576), False, 'from neuralcompression.functional import soft_round, soft_round_inverse\n'), ((775, 810), 'neuralcompression.functional.soft_round_inverse', 'soft_round_inverse', (['x'], {'alpha': '(5000.0)'}), '(x, alpha=5000.0)\n', (793, 810), False, 'from neuralcompression.functional import soft_round, soft_round_inverse\n'), ((742, 755), 'torch.ceil', 'torch.ceil', (['x'], {}), '(x)\n', (752, 755), False, 'import torch\n')]