Avoid state.aux=None in state returned by initial_state

Currently, when `has_aux=True`, `state.aux` is `None` when `state` is returned by `init_state` and `state.aux` is equal to `fun(params, *args, **kwargs)[1]` when `state` is returned by `update`. This is problematic as it can trigger a jit recompilation. One way would be to set `state.aux` to some dummy values of the correct type when returned by `init_state`.