tfm.core.task_factory.register_task_cls

Decorates a factory of Tasks for lookup by a subclass of TaskConfig.

tfm.core.task_factory.register_task_cls(
    task_config_cls
)

This decorator supports registration of tasks as follows:

@dataclasses.dataclass
class MyTaskConfig(TaskConfig):
  # Add fields here.
  pass

@register_task_cls(MyTaskConfig)
class MyTask(Task):
  # Inherits def __init__(self, task_config).
  pass

my_task_config = MyTaskConfig()
my_task = get_task(my_task_config)  # Returns MyTask(my_task_config).

Besisdes a class itself, other callables that create a Task from a TaskConfig can be decorated by the result of this function, as long as there is at most one registration for each config class.

task_config_cls a subclass of TaskConfig (not an instance of TaskConfig). Each task_config_cls can only be used for a single registration.

A callable for use as class decorator that registers the decorated class for creation from an instance of task_config_cls.