Skip to content

treex.Initializer

Initialize a field from a function that expects a single argument with a PRNGKey.

Initializers are called by Module.init and replace the value of the field they are assigned to.

__init__(self, f) special

Parameters:

Name Type Description Default
f Callable[[jax._src.numpy.lax_numpy.ndarray], Any]

A function that takes a PRNGKey and returns the initial value of the field.

required
Source code in treex/types.py
def __init__(self, f: tp.Callable[[jnp.ndarray], tp.Any]):
    """
    Arguments:
        f: A function that takes a PRNGKey and returns the initial value of the field.
    """
    self.f = f