Method initialize

axlearn/common/param_init.py:158–167

(
        self,
        name: str,
        *,
        prng_key: Tensor,
        shape: Shape,
        dtype: jnp.dtype,
        axes: Optional[FanAxes] = None,
    )

    """Base class for initializers."""

    def initialize(
        self,
        name: str,
        *,
        prng_key: Tensor,
        shape: Shape,
        dtype: jnp.dtype,
        axes: Optional[FanAxes] = None,
    ) -> jnp.ndarray:
        raise NotImplementedError(type(self))

    def debug_string(
        self,