def next_key(self) -> jnp.ndarray:
if LOCAL.rng is None:
raise types.NoContext(
f"Trying to call `next_key` from module {self} outside `init` or `apply`."
)
return LOCAL.rng.next()
def is_initializing(self) -> bool:
if LOCAL.initializing is None:
raise types.NoContext(
f"Trying to call `is_initializing` from module '{self}' outside `init` or `apply`."
)
return LOCAL.initializing
def _call_context(module: Module):
prev_parent = LOCAL.parent
prev_inside_call = LOCAL.inside_call
prev_module_index = LOCAL.module_index
LOCAL.parent = module
LOCAL.inside_call = True
LOCAL.module_index = 0
try:
if not has_scope():
raise types.NoContext(
f"Trying to call top-level module '{module}' directly, use `apply` instead."
)
elif prev_parent is not None and module not in prev_parent._submodule_name:
raise types.SubmoduleNotRegistered(
f"Submodule {utils.get_name(module)} not registered in {utils.get_name(prev_parent)}, "
f"this is probably due to some of the following reasons:\n"
f"- The submodule is being captured by closure and not registered to any field.\n"
f"- The submodule was set to a field of the parent but "
f"its contained inside a more complex type which elegy cannot "
f"inspect, elegy only looks structures of (possibly nested) list, tuple, or dict.\n"
f"- The submodule was set to a field of the parent by mutating such field after __init__\n\n"
f"- If non of the previous is true consider this a bug."
f"Submodule: {module}\n"
f"Module: {prev_parent}\n"
)
if prev_parent is None:
LOCAL.module_path[module] = ()
else:
parent_path = LOCAL.module_path[prev_parent]
child_name = prev_parent._submodule_name[module]
LOCAL.module_path[module] = parent_path + (child_name,)
yield
finally:
LOCAL.parent = prev_parent
LOCAL.inside_call = prev_inside_call
LOCAL.module_index = prev_module_index