Fix Glorot initialization, add He initialization

[Sleort]

Should fix #442 .
Adds He weight initialization as a bonus :-)

[MikeInnes]

Best to split this into two PRs to make it easier to review and merge.

[Sleort]

I couldn't figure out how to split the PR into two, so I deleted the He initialization stuff (for now).

(And sorry for the mess. I'm not yet completely familiar/comfortable with how Git(hub) works...)

[MikeInnes]

No worries. Splitting the commit is possible but it's usually easier just to delete code and/or copy-paste it somewhere else for the he PR, so that's the right thing to do.

[Sleort]

Anything else I should do to get this PR reviewed / merged?

[MikeInnes]

I'm a bit uncomfortable with the heuristic this uses, but at least it has good precedent in other libraries. We can always redesign in future if needed.

Thanks @Sleort!

bors r+

[bors]

Build succeeded

• ci/gitlab/gitlab.com

[Sleort]

I'm a bit uncomfortable with the heuristic this uses, but at least it has good precedent in other libraries. We can always redesign in future if needed.

I agree (see the discussion at #442). A more "proper" way of treating this would, however, require breaking changes, since we would have to dispatch on layer type in addition to dimensionality. Actually, ideally we should maybe make the initialization depend on the activation function as well, such that the "right" initialization could be automatically selected (Glorot for tanh, He for relu, etc.)...?

[MikeInnes]

We don't generally consider changing a layer's default initialisation to be breaking, so we could make the defaults smarter at any time, even if it isn't exposed as an API. A more general user-accessible initialisation API would be nice to have though.