class_docstr: ModuleBackend is a wrapper around torch module that allows it to run tasks asynchronously with Runtime
  By default, ModuleBackend handles three types of requests:
  
   - forward - receive inputs and compute outputs. Concurrent requests will be batched for better GPU utilization.
   - backward - receive gradients w.r.t. outputs, compute gradients w.r.t. inputs and **update expert**. Also batched.
   - get_info - return expert metadata. Not batched.
  
  :param module: nn.Module to be wrapped into a backend. Arbitrary pytorch module with a few limitations:
  
   - Experts must always receive the same set of args and kwargs and produce output tensors of same type
   - All args, kwargs and outputs must be **tensors** where 0-th dimension represents to batch size
   - We recommend using experts that are ~invariant to the order in which they process batches
   - Using randomness (e.g. Dropout) leads to different samples at forward and backward. If you want consistency,
      you should explicitly register these random variables as model inputs or outputs.
      See hivemind.utils.custom_layers.DeterministicDropout for an example
  
  :param optimizer: torch optimizer to be applied on every backward call
  :param scheduler: a function to create the learning rate scheduler for the expert
  :param args_schema: description of positional arguments to expert.forward, list of BatchTensorProto
  :param kwargs_schema: description of keyword arguments to expert.forward, dict of BatchTensorProto
  :param outputs_schema: description of outputs from expert.forward, nested structure of BatchTensorProto
  :param kwargs: extra parameters to be forwarded into TaskPool.__init__
class_docstr_tok: ["ModuleBackend","is","a","wrapper","around","torch","module","that","allows","it","to","run","tasks","asynchronously","with","Runtime","By","default",",","ModuleBackend","handles","three","types","of","requests",":","-","forward","-","receive","inputs","and","compute","outputs",".","Concurrent","requests","will","be","batched","for","better","GPU","utilization",".","-","backward","-","receive","gradients","w.r.t",".","outputs",",","compute","gradients","w.r.t",".","inputs","and","*","*","update","expert","*","*",".","Also","batched",".","-","get_info","-","return","expert","metadata",".","Not","batched",".",":","param","module",":","nn.Module","to","be","wrapped","into","a","backend",".","Arbitrary","pytorch","module","with","a","few","limitations",":","-","Experts","must","always","receive","the","same","set","of","args","and","kwargs","and","produce","output","tensors","of","same","type","-","All","args",",","kwargs","and","outputs","must","be","*","*","tensors","*","*","where","0-th","dimension","represents","to","batch","size","-","We","recommend","using","experts","that","are","~invariant","to","the","order","in","which","they","process","batches","-","Using","randomness","(","e.g",".","Dropout",")","leads","to","different","samples","at","forward","and","backward",".","If","you","want","consistency",",","you","should","explicitly","register","these","random","variables","as","model","inputs","or","outputs",".","See","hivemind.utils.custom_layers.DeterministicDropout","for","an","example",":","param","optimizer",":","torch","optimizer","to","be","applied","on","every","backward","call",":","param","scheduler",":","a","function","to","create","the","learning","rate","scheduler","for","the","expert",":","param","args_schema",":","description","of","positional","arguments","to","expert.forward",",","list","of","BatchTensorProto",":","param","kwargs_schema",":","description","of","keyword","arguments","to","expert.forward",",","dict","of","BatchTensorProto",":","param","outputs_schema",":","description","of","outputs","from","expert.forward",",","nested","structure","of","BatchTensorProto",":","param","kwargs",":","extra","parameters","to","be","forwarded","into","TaskPool.__init__"]
code: def __init__(
      self,
      name: str,
      module: nn.Module,
      *,
      optimizer: Optional[torch.optim.Optimizer] = None,
      scheduler: Optional[LRSchedulerBase] = None,
      args_schema: Tuple[BatchTensorDescriptor, ...] = None,
      kwargs_schema: Dict[str, BatchTensorDescriptor] = None,
      outputs_schema: Union[
          BatchTensorDescriptor, Tuple[BatchTensorDescriptor, ...]
      ] = None,
      **kwargs,
  ):
      super().__init__()
      self.name, self.module, self.optimizer, self.scheduler = (
          name,
          module,
          optimizer,
          scheduler,
      )
  
      self.args_schema = args_schema = tuple(args_schema or ())
      self.kwargs_schema = kwargs_schema = dict(kwargs_schema or {})
      assert args_schema or kwargs_schema, (
          f"Module must take at least one positional or keyword input."
          " Did you forget to provide args_schema/kwargs_schema?"
      )
      assert (
          optimizer is not None or scheduler is None
      ), "scheduler should only be used if optimizer is not None"
  
      if outputs_schema is None:
          # run expert once to get outputs schema
          dummy_args = tuple(
              sample.make_zeros(DUMMY_BATCH_SIZE)
              for sample in args_schema
          )
          dummy_kwargs = {
              key: sample.make_zeros(DUMMY_BATCH_SIZE)
              for key, sample in kwargs_schema.items()
          }
          dummy_outputs = self.module(*dummy_args, **dummy_kwargs)
          outputs_schema = nested_map(
              BatchTensorDescriptor.from_tensor, dummy_outputs
          )
  
      self.forward_schema = (
          self.args_schema,
          self.kwargs_schema,
      )  # inputs for forward
      self.outputs_schema = outputs_schema  # outputs from forward
  
      self.backward_schema = (
          self.forward_schema,
          self.outputs_schema,
      )  # inputs to backward
      self.grad_inputs_schema = (
          self.forward_schema
      )  # outputs from backward
      self.forward_pool = TaskPool(
          self.forward, name=f"{self.name}_forward", **kwargs
      )
      self.backward_pool = TaskPool(
          self.backward, name=f"{self.name}_backward", **kwargs
      )
  
file_path: hivemind/hivemind/moe/server/module_backend.py
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union
  import torch
  from torch import nn
  from hivemind.moe.server.task_pool import TaskPool
  from hivemind.utils.logging import get_logger
  from hivemind.utils.nested import nested_compare, nested_flatten, nested_map, nested_pack
  from hivemind.utils.tensor_descr import DUMMY_BATCH_SIZE, BatchTensorDescriptor
label_desc: Use image node_id 1 to create a new ModuleBackend object  with example: obj = ModuleBackend(name, module)