def register_task(
    name: str) -> Callable[[Type[_BaseTaskVar]], Type[_BaseTaskVar]]:
  """Register task.

  Task should implement BaseTask abstraction. Used as decorator, for example:

  @register_task('my_task')
  class MyTask(BaseTask):

  Args:
    name: name of registered task.

  Returns:
    Mapping from BaseTask to BaseTask.
  """

  def _wrap(cls: Type[_BaseTaskVar]) -> Type[_BaseTaskVar]:
    """Decorator inner wrapper needed to support `name` argument."""
    if not issubclass(cls, base_task.BaseTask):
      raise TypeError('Invalid task. Task %s does not subclass BaseTask.' %
                      cls.__name__)

    if name in _TASK_REGISTRY:
      raise ValueError(
          'Task name %s has already been registered with class %s' %
          (name, _TASK_REGISTRY[name].__name__))

    _TASK_REGISTRY[name] = cls

    return cls

  return _wrap
Register task.

  Task should implement BaseTask abstraction. Used as decorator, for example:

  @register_task('my_task')
  class MyTask(BaseTask):

  Args:
    name: name of registered task.

  Returns:
    Mapping from BaseTask to BaseTask.
  
language/mentionmemory/tasks/task_registry.py