code: def test_beam_search_correctness():
      all_expert_uids = [
          f"ffn.{5 + i}.{10 + j}.{15 + k}"
          for i in range(10)
          for j in range(10)
          for k in range(10)
      ]
      dht = DHT(start=True)
      assert all(
          declare_experts(
              dht, all_expert_uids, expiration_time=get_dht_time() + 30
          )
      )
  
      dmoe = RemoteMixtureOfExperts(
          in_features=32,
          grid_size=(32, 32, 32),
          dht=dht,
          k_best=4,
          uid_prefix="ffn.",
      )
  
      for _ in range(25):
          input = torch.randn(32)
          grid_scores = dmoe.proj(input).split_with_sizes(
              dmoe.beam_search.grid_size, dim=-1
          )
  
          chosen_experts = dmoe.beam_search.find_best_experts(
              [tensor.detach().numpy() for tensor in grid_scores],
              beam_size=dmoe.k_best,
          )
          chosen_scores = dmoe.compute_expert_scores(
              [dim_scores[None] for dim_scores in grid_scores],
              [chosen_experts],
          )[0]
          our_best_scores = list(chosen_scores.cpu().detach().numpy())
  
          # reference: independently find :beam_size: best experts with exhaustive search
          all_scores = dmoe.compute_expert_scores(
              [dim_scores.unsqueeze(0) for dim_scores in grid_scores],
              [
                  [
                      RemoteExpert(ExpertInfo(uid, None), None)
                      for uid in all_expert_uids
                  ]
              ],
          )[0]
          true_best_scores = sorted(
              all_scores.cpu().detach().numpy(), reverse=True
          )[: len(chosen_experts)]
  
          assert np.allclose(true_best_scores, our_best_scores)
  
imports: import asyncio
  import ctypes
  import multiprocessing
  import threading
  import time
  import numpy
  import pytest
  import torch
  from hivemind.dht import DHT
  from hivemind.moe.client.expert import RemoteExpert, create_remote_experts
  from hivemind.moe.client.moe import DUMMY, RemoteMixtureOfExperts, _RemoteCallMany
  from hivemind.moe.client.remote_expert_worker import RemoteExpertWorker
  from hivemind.moe.client.switch_moe import RemoteSwitchMixtureOfExperts
  from hivemind.moe.expert_uid import ExpertInfo
  from hivemind.moe.server import ModuleBackend, Server, background_server, declare_experts
  from hivemind.moe.server.layers import name_to_block
  from hivemind.p2p.p2p_daemon_bindings.control import P2PHandlerError
  from hivemind.utils import BatchTensorDescriptor, MPFuture, get_dht_time
label_desc: Use image node_id 5 for calling a global function with example usage: test_beam_search_correctness() without return types