Allow complex values in Linear layers

penzai/nn/linear_and_affine.py

@@ -586,7 +586,7 @@ def _input_structure(self):
     }
     return shapecheck.ArraySpec(
         named_shape={**shapecheck.var("B"), **known_in_axes},
-        dtype=jnp.floating,
+        dtype=jnp.inexact,
     )
 
   def _output_structure(self):
@@ -597,7 +597,7 @@ def _output_structure(self):
     }
     return shapecheck.ArraySpec(
         named_shape={**shapecheck.var("B"), **known_out_axes},
-        dtype=jnp.floating,
+        dtype=jnp.inexact,
     )
 
   def treescope_color(self) -> str:

tests/nn/linear_and_affine_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 import chex
 import jax
+import jax.numpy as jnp
 from penzai import pz
 
 
@@ -108,6 +109,25 @@ def test_linear_not_in_place(self):
         ),
     )
 
+  def test_linear_supports_complex_values(self):
+    layer = pz.nn.Linear.from_config(
+        name="complex_linear",
+        init_base_rng=jax.random.key(1),
+        input_axes={"foo": 3},
+        output_axes={"bar": 5},
+        dtype=jnp.complex64,
+        rename_outputs_if_necessary=False,
+    )
+    result = layer(
+        pz.nx.ones({"batch": 2, "foo": 3}, dtype=jnp.complex64) * (1 + 2j)
+    )
+    pz.chk.check_structure(
+        result,
+        pz.chk.ArraySpec(
+            named_shape={"batch": 2, "bar": 5}, dtype=jnp.complex64
+        ),
+    )
+
   def test_linear_in_place(self):
     layer = pz.nn.Linear.from_config(
         name="test",