From 31d4af921e19b8b6518af8586c71bb82b8561c18 Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Sat, 15 Aug 2026 14:10:22 +0100
Subject: [PATCH] Allow complex values in Linear layers

---
 penzai/nn/linear_and_affine.py     |  4 ++--
 tests/nn/linear_and_affine_test.py | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/penzai/nn/linear_and_affine.py b/penzai/nn/linear_and_affine.py
index f696faf..ba68262 100644
--- a/penzai/nn/linear_and_affine.py
+++ b/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:
diff --git a/tests/nn/linear_and_affine_test.py b/tests/nn/linear_and_affine_test.py
index 3c62467..ba2bb39 100644
--- a/tests/nn/linear_and_affine_test.py
+++ b/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",