# 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.

Source code in treex/types.py
class Initializer(to.Tree):
"""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.
"""

f: tp.Callable[[jnp.ndarray], tp.Any]

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

def __call__(self, x: jnp.ndarray) -> np.ndarray:
return self.f(x)

def __repr__(self) -> str:
return "Initializer"


## __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