treex.rng_key

Source code in treex/module.py

@contextmanager
def rng_key(key: types.KeyLike):
    key = utils.Key(key)

    with _INIT_CONTEXT.update(key=key):
        yield