Support complex dtypes in networks#765

Merged

patrick-kidger merged 5 commits into

patrick-kidger:mainfrom

ChenAo-Phys:main

Jun 24, 2024

Commits on Jun 21, 2024

add support for complex networks
ChenAo-Phys
committed
add tests for complex networks
ChenAo-Phys
committed
support for older version of jax
ChenAo-Phys
committed
replace PositionalSharding by NamedSharding
ChenAo-Phys
committed

Commits on Jun 22, 2024

move default_init to nn._misc
ChenAo-Phys
committed