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",