def apply_recursive(context: tp.Tuple[str, ...], metrics, **kwargs):
if isinstance(metrics, tp.Callable):
name = (
metrics.name
if isinstance(metrics, module.Module)
else utils.lower_snake_case(metrics.__name__)
)
context += (name,)
value = utils.inject_dependencies(metrics)(**kwargs)
if isinstance(value, tp.Dict):
for name, value in value.items():
yield context + (name,), value
else:
yield context, value
elif isinstance(metrics, (tp.Tuple, tp.List)):
for loss in metrics:
yield from apply_recursive(context, loss, **kwargs)
elif isinstance(metrics, tp.Dict):
for name, loss in metrics.items():
yield from apply_recursive(context + (name,), loss, **kwargs)
else:
raise TypeError(f"Invalid type {type(metrics)}")
def __init__(self,
name: tp.Optional[str] = None,
dtype: tp.Any = jnp.float32):
"""
Initializes the current module with the given name.
Subclasses should call this constructor before creating other modules or
variables such that those modules are named correctly.
Arguments:
name: An optional string name for the class. If ``name`` is not provided then the class name for the
current instance is converted to ``lower_snake_case`` and used instead.
"""
self.name = name if name else utils.lower_snake_case(
self.__class__.__name__)
self.dtype = dtype
self._submodules = {}
self._submodule_name = {}
self._dynamic_submodules = []
self._default_params = {}
self._scope_params = None
self._spec = {}
self._initialized = False
self._signature_f = self.call
self._jit_functions()
def apply_recursive(self, context: tp.Tuple[str, ...], losses, **kwargs):
if isinstance(losses, tp.Callable):
name = (
losses.name
if isinstance(losses, Loss)
else utils.lower_snake_case(losses.__name__)
)
context += (name,)
val = utils.inject_dependencies(losses)(**kwargs)
if isinstance(val, tp.Dict):
for name, val in val.items():
yield context + (name,), val
else:
yield context, val
elif isinstance(losses, (tp.Tuple, tp.List)):
for loss in losses:
yield from self.apply_recursive(context, loss, **kwargs)
elif isinstance(losses, tp.Dict):
for name, loss in losses.items():
yield from self.apply_recursive(context + (name,), loss, **kwargs)
else:
raise TypeError(f"Invalid type {type(losses)}")
def __init__(self,
name: tp.Optional[str] = None,
dtype: np.dtype = jnp.float32):
"""
Initializes the current module with the given name.
Subclasses should call this constructor before creating other modules or
variables such that those modules are named correctly.
Arguments:
name: An optional string name for the class. Must be a valid elsePython
identifier. If ``name`` is not provided then the class name for the
current instance is converted to ``lower_snake_case`` and used instead.
"""
self.name = name if name else utils.lower_snake_case(
self.__class__.__name__)
self.dtype = dtype
self._params = []
self._states = []
self._submodules = []
self._dynamic_submodules = []
self._ignore = []
def __init__(self, name: tp.Optional[str] = None):
super().__init__(name=utils.lower_snake_case(f.__name__)
if name is None else name)
self.call = f