small fixes
small fixes

- Now only supports the unbatched case
- Changes initialization computation to use jnp.prod instead of math.prod
- Enforces correct tuple length in parse() 
- Doesn't restate fields in inherited classes
- Some style fixes
patrick-kidger/equinox