I’m super interested in FP for AGI, but one issue with such setups (see Jax.experimental.stax) is the neural nets are parametric functions, and if recurrent, ca

I’m super interested in FP for AGI, but one issue with such setups (see Jax.experimental.stax) is the neural nets are parametric functions, and if recurrent, can have state, too. That adds up to a lot of extra work for a programmer to manually pass around these parameters. It’s doable but also distracting from the real work of cognitive architecture, which is why I switched from stax style to Flax. Flax is currently by far my favorite neural net library, despite being pretty small, because it lets you write a neural net module as a function and decorate that with @nn.module — the decorator wraps your function in a class instance to manage the params and state. This can shave 10 lines per module, and if you’re like me, you have like 50 bajillion ideas for different neural net modules, it really adds up. TLDR: Check out Flax if you like functional programming and neural networks!