class_docstr: This factory is used to create task models bundling the task,
     encoder, and task head within the task model.
  
  Attributes:
      registry (dict): Dynamic registry mapping task types to task models
class_docstr_tok: ["This","factory","is","used","to","create","task","models","bundling","the","task",",","encoder",",","and","task","head","within","the","task","model",".","Attributes",":","registry","(","dict",")",":","Dynamic","registry","mapping","task","types","to","task","models"]
docstr: Register task_type as a key mapping to a TaskModel
  
  Args:
      task_type (TaskTypes): TaskType key mapping to a BaseHead task head
  
  Returns:
      Callable: inner_wrapper() wrapping TaskModel constructor
docstr_tok: ["Register","task_type","as","a","key","mapping","to","a","TaskModel","Args",":","task_type","(","TaskTypes",")",":","TaskType","key","mapping","to","a","BaseHead","task","head","Returns",":","Callable",":","inner_wrapper","(",")","wrapping","TaskModel","constructor"]
code: def register(cls, task_type: TaskTypes) -> Callable:
      """Register task_type as a key mapping to a TaskModel
  
      Args:
          task_type (TaskTypes): TaskType key mapping to a BaseHead task head
  
      Returns:
          Callable: inner_wrapper() wrapping TaskModel constructor
      """
  
      def inner_wrapper(wrapped_class: Taskmodel) -> Callable:
          assert task_type not in cls.registry
          cls.registry[task_type] = wrapped_class
          return wrapped_class
  
      return inner_wrapper
  
imports: import abc
  import torch
  import torch.nn
  from typing import Callable
  import jiant.proj.main.modeling.heads
  from jiant.proj.main.components.outputs import LogitsAndLossOutput
  from jiant.proj.main.components.outputs import LogitsOutput
  from jiant.utils.python.datastructures import take_one
  from jiant.tasks.core import TaskTypes
label_desc: Use image node_id 1 for calling the JiantTaskModelFactory obj's underlying member method code with example usage: obj.register(cls, task_type) and returns: inner_wrapper, wrapped_class