Skip to content
This repository was archived by the owner on Aug 25, 2026. It is now read-only.

Allow complex values in Linear layers - #141

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix-linear-complex-dtypes
Open

Allow complex values in Linear layers#141
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix-linear-complex-dtypes

Conversation

@sylvesterkaczmarek

@sylvesterkaczmarek sylvesterkaczmarek commented Aug 15, 2026

Copy link
Copy Markdown

Fixes #127.

Summary

Allow Linear layers to accept complex-valued inputs and outputs by using NumPy/JAX's inexact abstract dtype rather than restricting structure checks to floating-point values.

  • change the input/output ArraySpec dtype from jnp.floating to jnp.inexact
  • add a complex64 Linear regression test

Why

JAX linear algebra supports complex values, but Penzai's structure declaration rejected them before the layer could run. jnp.inexact is the abstract dtype family covering both floating and complex values.

Testing

Targeted linear/affine tests passed on the branch: 12 tests. Formatting checks also passed.

@sylvesterkaczmarek
sylvesterkaczmarek marked this pull request as ready for review August 15, 2026 16:55
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Using Complex Numbers in the Linear layer does not work.

1 participant