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 backward(
      cls, ctx, *raw_grads: torch.Tensor
  ) -> Tuple[Optional[torch.Tensor], ...]:
      assert not torch.is_grad_enabled()
      (
          info,
          backward_k_min,
          backward_timeout,
          timeout_after_k_min,
          expert_per_sample,
          detect_anomalies,
      ) = ctx._saved_non_tensors
      alive_ii, alive_jj, *flat_inputs_cpu = ctx.saved_tensors
  
      dummy_grad_mask, *flat_grad_outputs = raw_grads
  
      flat_grad_outputs_cpu = []
      for tensor in flat_grad_outputs:
          if detect_anomalies and not tensor.isfinite().all():
              raise ValueError("One of gradients has nan/inf values")
          flat_grad_outputs_cpu.append(tensor.cpu())
  
      num_samples, max_experts = dummy_grad_mask.shape
  
      inputs_per_expert = zip(
          *(
              tensor[alive_ii].split(1, dim=0)
              for tensor in flat_inputs_cpu
          )
      )
      grad_outputs_per_expert = zip(
          *(
              tensor[alive_ii, alive_jj].split(1, dim=0)
              for tensor in flat_grad_outputs_cpu
          )
      )
      backward_schema = tuple(
          nested_flatten(
              (info["forward_schema"], info["outputs_schema"])
          )
      )
  
      # dispatch tasks to all remote experts, collect responses
      pending_tasks = {}
      for i, j, inputs_ij, grad_outputs_ij in zip(
          alive_ii.cpu().numpy(),
          alive_jj.cpu().numpy(),
          inputs_per_expert,
          grad_outputs_per_expert,
      ):
          expert: RemoteExpert = expert_per_sample[i.item()][j.item()]
          stub = get_server_stub(expert.p2p, expert.peer_id)
          inputs_and_grad_outputs = tuple(
              nested_flatten((inputs_ij, grad_outputs_ij))
          )
          serialized_tensors = (
              serialize_torch_tensor(tensor, proto.compression)
              for tensor, proto in zip(
                  inputs_and_grad_outputs, backward_schema
              )
          )
          new_task = RemoteExpertWorker.run_coroutine(
              expert_backward(
                  expert.uid,
                  inputs_and_grad_outputs,
                  serialized_tensors,
                  stub,
              ),
              return_future=True,
          )
          pending_tasks[new_task] = (i, j)
  
      survivor_inds, survivor_grad_inputs = cls._collect_responses(
          pending_tasks,
          num_samples,
          backward_k_min,
          backward_timeout,
          timeout_after_k_min,
          detect_anomalies,
      )
      if len(survivor_inds) < backward_k_min:
          raise TimeoutError(
              f"Backward pass: less than {backward_k_min} experts responded within timeout"
          )
  
      # assemble responses
      batch_inds, expert_inds = map(
          lambda x: torch.as_tensor(x, dtype=torch.long),
          list(zip(*survivor_inds)) or ([], []),
      )
  
      survivor_grad_inputs_stacked = (
          torch.cat(grad_inputs)
          for grad_inputs in zip(*survivor_grad_inputs)
      )
      # torch tensors, i-th tensor is of shape [num_backward_survivors, *flat_inputs_cpu[i].shape]
  
      grad_inputs = nested_map(
          lambda descr: descr.make_zeros(
              num_samples, device=flat_grad_outputs[0].device
          ),
          list(nested_flatten(info["forward_schema"])),
      )
  
      for grad_input, survivor_grad_stacked in zip(
          grad_inputs, survivor_grad_inputs_stacked
      ):
          grad_input_per_expert = torch.zeros(  # gradient tensor with individual contributions from each expert
              (num_samples, max_experts, *grad_input.shape[1:]),
              device=survivor_grad_stacked.device,
              dtype=survivor_grad_stacked.dtype,
          )
          grad_input_per_expert[
              batch_inds, expert_inds
          ] = survivor_grad_stacked
          grad_input.copy_(
              grad_input_per_expert.to(flat_grad_outputs[0].device).sum(
                  dim=1
              )
          )
  
      return (
          DUMMY,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          *grad_inputs,
      )
  
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 2 for calling the _RemoteCallMany obj's underlying member method code with example usage: obj.backward(cls, ctx) and returns: DUMMY,  None,  None,  None,  None,  None,  None,  None,  None,  None,  unknown