Allow zero shapes in nn.Linear

Author: aseyboldt

equinox/nn/_linear.py

@@ -52,7 +52,10 @@ def __init__(
         wkey, bkey = jrandom.split(key, 2)
         in_features_ = 1 if in_features == "scalar" else in_features
         out_features_ = 1 if out_features == "scalar" else out_features
-        lim = 1 / math.sqrt(in_features_)
+        if in_features_ == 0:
+            lim = 1.0
+        else:
+            lim = 1 / math.sqrt(in_features_)
         wshape = (out_features_, in_features_)
         self.weight = default_init(wkey, wshape, dtype, lim)
         bshape = (out_features_,)

tests/test_nn.py

@@ -22,6 +22,16 @@ def test_custom_init():
 
 
 def test_linear(getkey):
+    # Zero input shape
+    linear = eqx.nn.Linear(0, 4, key=getkey())
+    x = jrandom.normal(getkey(), (0,))
+    assert linear(x).shape == (4,)
+
+    # Zero output shape
+    linear = eqx.nn.Linear(4, 0, key=getkey())
+    x = jrandom.normal(getkey(), (4,))
+    assert linear(x).shape == (0,)
+
     # Positional arguments
     linear = eqx.nn.Linear(3, 4, key=getkey())
     x = jrandom.normal(getkey(), (3,))