Function bind_module

axlearn/common/test_utils.py:776–827

Creates a context in which `module` has `state`.` This lets you write tests that make calls to a module without needing to call `functional()` yourself. It is similar in spirit to FLAX's `module.bind()` although that works differently due to the fact that FLAX state is only associa

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

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

    This lets you write tests that make calls to a module without needing to call `functional()`
    yourself.

    It is similar in spirit to FLAX's `module.bind()` although that works differently due to the
    fact that FLAX state is only associated with an instance of a module, whereas AXLearn state is
    global.

    Example:
        ```
        cfg = MyModule.default_config()
        with test_utils.bind_layer(cfg) as module:
            result = module.do_something(some_args)
        ```

    Args:
        module: The module to create a context for.
        is_training: Tell the module it is in training or not.
        prng_key: The PRNG key to use. If None, `jax.random.PRNGKey(0)`.
        state: The state to use.

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

    if isinstance(module, InstantiableConfig):
        if isinstance(module, Module.Config) and isinstance(
            getattr(module, "name", None), RequiredFieldValue
        ):
            setattr(module, "name", "tmp")
        module = module.instantiate(parent=None)
    ctx = InvocationContext(
        name="root",
        parent=None,
        module=module,
        is_training=is_training,
        prng_key=prng_key,
        state=state,
        output_collection=new_output_collection(),
    )
    with set_current_context(ctx):
        yield module


@contextlib.contextmanager