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

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
@sylvesterkaczmarek
sylvesterkaczmarek force-pushed the fix-linear-complex-dtypes branch from 31d4af9 to 6851ab5 Compare October 1, 2026 19:42
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

1 participant