fix Initializer protocol

[Qwlouse]

The initializer protocol is currently defined in terms of a static __call__ method.
However that definition does not interact correctly with callable classes. For example this raises a mypy error (but no pytype error):

import jax

class MyInit:
  def __call__(self, key, shape, dtype=None):
    pass

def foo(init: jax.nn.Initializer):
  return

foo(MyInit())

After this PR the above code passes the mypy check, and passing plain functions still works as before:

import jax

def my_init(key, shape, dtype=None):
  return

def foo(init: jax.nn.Initializer):
  return

foo(my_init)