class_docstr: Internal autograd-friendly function that calls multiple experts on a batch of inputs and awaits responses
  This function that can recover from individual failures during forward and/or backward pass as long as at least
  one expert succeeds for each input. For user-friendly version of this function, use RemoteMixtureOfExperts module.
  
  Note: experts that failed during forward will be assigned zero outputs and marked as mask[i, j] = 0,
        experts that failed during backward will be treated as constants (i.e. gradients through them are zeros)
class_docstr_tok: ["Internal","autograd-friendly","function","that","calls","multiple","experts","on","a","batch","of","inputs","and","awaits","responses","This","function","that","can","recover","from","individual","failures","during","forward","and\/or","backward","pass","as","long","as","at","least","one","expert","succeeds","for","each","input",".","For","user-friendly","version","of","this","function",",","use","RemoteMixtureOfExperts","module",".","Note",":","experts","that","failed","during","forward","will","be","assigned","zero","outputs","and","marked","as","mask","[","i",",","j","]","=","0",",","experts","that","failed","during","backward","will","be","treated","as","constants","(","i.e",".","gradients","through","them","are","zeros",")"]
code: def forward(
      cls,
      ctx,
      dummy,
      experts_per_sample: List[List[RemoteExpert]],
      k_min: int,
      backward_k_min: int,
      timeout_after_k_min: float,
      forward_timeout: Optional[float],
      backward_timeout: Optional[float],
      detect_anomalies: bool,
      allow_zero_outputs: bool,
      info: Dict[str, Any],
      *flat_inputs: torch.Tensor,
  ) -> Tuple[torch.Tensor]:
      assert not torch.is_grad_enabled()
      num_samples, max_experts = len(experts_per_sample), max(
          map(len, experts_per_sample)
      )
  
      flat_inputs_cpu = []
      for tensor in flat_inputs:
          if detect_anomalies and not tensor.isfinite().all():
              raise ValueError("One of inputs has nan/inf values")
          flat_inputs_cpu.append(tensor.cpu())
  
      flat_inputs_per_sample = list(
          zip(*(x.split(1, dim=0) for x in flat_inputs_cpu))
      )
      assert (
          len(experts_per_sample)
          == len(flat_inputs_per_sample)
          == num_samples
      )
  
      # dispatch tasks to all remote experts collect responses
      pending_tasks: Dict[Future, Tuple[int, int]] = {}
      for i in range(num_samples):
          for j, expert in enumerate(experts_per_sample[i]):
              stub = get_server_stub(expert.p2p, expert.peer_id)
              serialized_tensors = (
                  serialize_torch_tensor(tensor, proto.compression)
                  for tensor, proto in zip(
                      flat_inputs_per_sample[i],
                      nested_flatten(info["forward_schema"]),
                  )
              )
              new_task = RemoteExpertWorker.run_coroutine(
                  expert_forward(
                      expert.uid,
                      flat_inputs_per_sample[i],
                      serialized_tensors,
                      stub,
                  ),
                  return_future=True,
              )
              pending_tasks[new_task] = (i, j)
  
      responded_inds, alive_flat_outputs = cls._collect_responses(
          pending_tasks,
          num_samples,
          k_min,
          forward_timeout,
          timeout_after_k_min,
          detect_anomalies,
      )
      if len(responded_inds) < k_min:
          raise TimeoutError(
              f"Forward pass: less than {k_min} responded within timeout"
          )
  
      if not isinstance(info["outputs_schema"], tuple):
          outputs_schema = (info["outputs_schema"],)
      else:
          outputs_schema = info["outputs_schema"]
      outputs = nested_map(
          lambda descriptor: descriptor.make_zeros(
              num_samples, max_experts, device=flat_inputs[0].device
          ),
          outputs_schema,
      )
  
      # assemble responses
      if len(responded_inds) > 0 or allow_zero_outputs:
          batch_inds, expert_inds = map(
              lambda x: torch.as_tensor(
                  x, device=flat_inputs[0].device, dtype=torch.long
              ),
              list(zip(*responded_inds)) or ([], []),
          )
  
          alive_flat_outputs_stacked = (
              torch.cat(outputs) for outputs in zip(*alive_flat_outputs)
          )
          # torch tensors, i-th tensor is of shape [num_responded, *expert_outputs[i].shape]
  
          for output, response_stacked in zip(
              outputs, alive_flat_outputs_stacked
          ):
              output[batch_inds, expert_inds] = response_stacked.to(
                  output.device
              )
  
      else:
          raise RuntimeError(
              "Forward pass: 0 experts responded within timeout and allow_zero_outputs is False"
          )
  
      mask = torch.zeros(
          [num_samples, max_experts],
          dtype=torch.bool,
          device=flat_inputs[0].device,
      )
      mask[batch_inds, expert_inds] = True
  
      # save individual outputs for backward pass
      ctx.save_for_backward(batch_inds, expert_inds, *flat_inputs_cpu)
      ctx._saved_non_tensors = (
          info,
          backward_k_min,
          backward_timeout,
          timeout_after_k_min,
          experts_per_sample,
          detect_anomalies,
      )
  
      return (mask,) + outputs
  
imports: from __future__ import annotations
  import time
  from concurrent.futures import Future
  from queue import Empty, Queue
  from typing import Any, Dict, List, Optional, Tuple
  import torch
  import torch.nn
  from torch.autograd.function import once_differentiable
  from hivemind.compression import serialize_torch_tensor
  from hivemind.dht import DHT
  from hivemind.moe.client.beam_search import MoEBeamSearcher
  from hivemind.moe.client.expert import DUMMY, RemoteExpert, expert_backward, expert_forward, get_server_stub
  from hivemind.moe.client.remote_expert_worker import RemoteExpertWorker
  from hivemind.moe.expert_uid import UID_DELIMITER
  from hivemind.p2p.p2p_daemon_bindings.control import P2PDaemonError
  from hivemind.utils import nested_flatten, nested_map, nested_pack
  from hivemind.utils.logging import get_logger
label_desc: Use image node_id 1 for calling the _RemoteCallMany obj's underlying member method code with example usage: obj.forward(cls, ctx, dummy, experts_per_sample, k_min, backward_k_min, timeout_after_k_min, forward_timeout, backward_timeout, detect_anomalies, allow_zero_outputs, info) and returns: unknown