code: def test_scatter(n_expert, topk, batch_size, d_model, world_size):
      gate_idx = torch.randint(n_expert + 1, (batch_size, topk)) - 1
      gate_idx = gate_idx.long().cuda()
      pos, lec, gec = count_by_gate(gate_idx, n_expert, world_size)
      fbs = int(gec.sum().item())
      inp = torch.rand(batch_size, d_model).cuda()
      inp.requires_grad = True
      out = MOEScatter.apply(
          inp, pos % batch_size, lec, gec, fbs, world_size
      )
      out.sum().backward()
  
      inp_raw = inp.data.clone()
      out_raw = torch.empty(
          pos.shape[0], d_model, device=inp.device, dtype=inp.dtype
      )
      # out_raw.sum().backward()
      for i, f in enumerate(pos.cpu()):
          out_raw[i] = inp[f % batch_size]
      _assert_numerical(["out"], [out], [out_raw], 0)
  
file_path: fastmoe/tests/test_local_exchange.py
imports: import sys
  from collections import OrderedDict
  from typing import List, Type, Union
  import pytest
  import torch
  import torch.nn
  import numpy
  from copy import deepcopy
  from fmoe.functions import MOEGather, MOEScatter, count_by_gate
  from test_numerical import _assert_numerical
label_desc: Use image node_id 1 for calling a global function with example usage: test_scatter(n_expert, topk, batch_size, d_model, world_size) without return types