We should support passing in custom initialization functions for weight initializations.
We should support passing in custom initialization functions for weight initializations.