code: def check_all_set_all_get_optm_state(
      device, state_step, state_mem, state_power
  ):
      num_embs = state_mem.shape[0]
      emb_dim = state_mem.shape[1]
      dgl_emb = NodeEmbedding(num_embs, emb_dim, "test", device=device)
      optm = SparseAdam(params=[dgl_emb], lr=0.01)
  
      dgl_emb._all_set_optm_state((state_step, state_mem, state_power))
  
      out_step, out_mem, out_power = dgl_emb._all_get_optm_state()
  
      assert F.allclose(state_step, out_step)
      assert F.allclose(state_mem, out_mem)
      assert F.allclose(state_power, out_power)
  
file_path: dgl/tests/python/pytorch/nn/test_sparse_emb.py
imports: import multiprocessing
  import os
  import unittest
  import backend
  import pytest
  import torch
  from dgl.nn import NodeEmbedding
  from dgl.optim import SparseAdam
label_desc: Use image node_id 3 for calling a global function with example usage: check_all_set_all_get_optm_state(device, state_step, state_mem, state_power) without return types