Function bind_layer

axlearn/common/test_utils.py:831–871

Creates a context in which `layer` has state initialized using `initialize_parameters_recursively`. The only difference between this and `bind_module()` is this calls `initialize_parameters_recursively`. Example: ``` cfg = Linear.default_config().set(input_dim=5, ou

(
    layer: ConfigOr[L],
    *,
    is_training: bool = True,
    prng_key: Optional[jax.random.PRNGKey] = None,
    state: Optional[Nested[Tensor]] = None,
)

@contextlib.contextmanager
def bind_layer(
    layer: ConfigOr[L],
    *,
    is_training: bool = True,
    prng_key: Optional[jax.random.PRNGKey] = None,
    state: Optional[Nested[Tensor]] = None,
) -> Iterator[L]:
    """Creates a context in which `layer` has state initialized using
    `initialize_parameters_recursively`.

    The only difference between this and `bind_module()` is this calls
    `initialize_parameters_recursively`.

    Example:
        ```
        cfg = Linear.default_config().set(input_dim=5, output_dim=7)
        with test_utils.bind_layer(cfg) as layer:
            result = layer(jnp.ones(5))
        assert result.shape == (7,)
        ```

    Args:
        layer: The layer to initialize.
        is_training: Tell the layer it is in training or not.
        prng_key: The PRNG key to use. If None, `jax.random.PRNGKey(0)`.
        state: The state to use. If None, call `initialize_parameters_recursively()` to initialize
               the state.

    Returns:
        The Initialized module.
    """
    if prng_key is None:
        prng_key = jax.random.PRNGKey(0)

    init_key, ctx_key = jax.random.split(prng_key)

    with bind_module(layer, is_training=is_training, prng_key=ctx_key, state={}) as instance:
        if state is None:
            state = instance.initialize_parameters_recursively(prng_key=init_key)
        current_context().state = state
        yield instance


def initialize_parameters_with_prebuilt(