From 3e6574b168c83ad6b68b469239d5e52a10e445ba Mon Sep 17 00:00:00 2001 From: Pablosinyores Date: Wed, 1 Jul 2026 17:01:40 +0530 Subject: [PATCH] Fix HardShrink to accept its documented lambd argument HardShrink was built with @_make_activation_module, which sets __call__ to always use the default lambd and gives the class no __init__. Its docstring documents a configurable lambd, but the class could not be constructed with one: nn.HardShrink(lambd=0.3) raised TypeError. Replace the decorator with an explicit __init__/__call__ that stores and forwards lambd, mirroring the sibling Softshrink layer. Extend test_hard_shrink to exercise the class path (default and lambd=0.1), which was previously untested. --- python/mlx/nn/layers/activations.py | 8 +++++++- python/tests/test_nn.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/activations.py b/python/mlx/nn/layers/activations.py index 9249499238..4f36ae2b38 100644 --- a/python/mlx/nn/layers/activations.py +++ b/python/mlx/nn/layers/activations.py @@ -642,7 +642,6 @@ class HardTanh(Module): """ -@_make_activation_module(hard_shrink) class HardShrink(Module): r"""Applies the HardShrink function. @@ -652,6 +651,13 @@ class HardShrink(Module): lambd: the :math:`\lambda` value for Hardshrink. Default: ``0.5`` """ + def __init__(self, lambd=0.5): + super().__init__() + self.lambd = lambd + + def __call__(self, x): + return hard_shrink(x, self.lambd) + @_make_activation_module(softmin) class Softmin(Module): diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 79974aac48..8ad72a323c 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -1185,6 +1185,18 @@ def test_hard_shrink(self): self.assertEqual(y.shape, (5,)) self.assertEqual(y.dtype, mx.float32) + y = nn.HardShrink()(x) + expected_y = mx.array([1.0, 0.0, 0.0, 0.0, -1.5]) + self.assertTrue(mx.array_equal(y, expected_y)) + self.assertEqual(y.shape, (5,)) + self.assertEqual(y.dtype, mx.float32) + + y = nn.HardShrink(lambd=0.1)(x) + expected_y = mx.array([1.0, -0.5, 0.0, 0.5, -1.5]) + self.assertTrue(mx.array_equal(y, expected_y)) + self.assertEqual(y.shape, (5,)) + self.assertEqual(y.dtype, mx.float32) + def test_rope(self): for kwargs in [{}, {"traditional": False}, {"base": 10000}, {"scale": 0.25}]: rope = nn.RoPE(4, **kwargs)