Fix HardShrink to accept its documented lambd argument

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):

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)