WIP] Zephyr NN Jax FrameworkL Short, Simple, Declarative

Hacker News

Share card

WIP] Zephyr NN Jax FrameworkL Short, Simple, Declarative

Hello HN! I have an early work-in-progress Neural Network Framework written on top of JAX. Simple | Declarative | No need to learn duplicated JAX transforms or specialized manipulation functions Its key difference with other JAX frameworks is its simplicity and straightforwardness. With other frameworks, the network is transformed to an (init, apply) which are pure functions and you basically don't use the actual code you've written and instead use this. With zephyr, neural networks look like neural networks: a function of parameters, input-data, and hyperparameters. It's also patterned for FP use, so partial application will a useful alternative to states in OO. Lastly, it's meant to be declarative and simple. Neural networks are just functions, not objects that need instantiation or anything. This means code are usually shorter as declaration and usage/computation happens in the same place since those things are highly coupled and so placing them together results in less cognitive load. It's early stage and so the core nets are few and unpolished, but I want to focus on the core first before moving on to implementing all core nets. Feedback on any part of it is very welcome!