code: def test_cpu_adagrad_opt_sparse_embedding(
      self, model_size=32, vocabulary_size=64, dim=16
  ):
      device = "cpu"
      rng_state = torch.get_rng_state()
  
      def gen_sparse_grad(
          vocabulary_size, dim, num_indices, dtype, device
      ):
          i = torch.randint(
              vocabulary_size,
              size=(1, num_indices),
              dtype=torch.int64,
              device=device,
          )
          v = torch.randn(num_indices, dim, dtype=dtype, device=device)
          t = torch.sparse_coo_tensor(
              i, v, (vocabulary_size, dim), device=device
          )
          t = t.coalesce()
          new_i = (
              (
                  t.indices().view(-1, 1).repeat(1, dim) * dim
                  + torch.tensor(range(dim))
              )
              .flatten()
              .unsqueeze(0)
          )
          new_v = t.values().flatten()
          new_t = torch.sparse_coo_tensor(
              new_i, new_v, (vocabulary_size * dim,), device=device
          )
          new_t = new_t.coalesce()
          new_t.requires_grad = False
          return new_t
  
      voc_size = vocabulary_size
      dim = dim
      num_indices = int(model_size // dim)
      dtype = torch.float32
  
      param = torch.nn.Parameter(
          torch.randn((voc_size * dim,), dtype=dtype, device=device),
          requires_grad=True,
      )
      torch.set_rng_state(rng_state)
      param1 = torch.nn.Parameter(
          torch.randn((voc_size * dim,), dtype=dtype, device=device),
          requires_grad=True,
      )
      torch.set_rng_state(rng_state)
  
      optimizer = DeepSpeedCPUAdagrad([param])
      optimizer1 = torch.optim.Adagrad([param1])
  
      for i in range(10):
          torch.set_rng_state(rng_state)
          param.grad = gen_sparse_grad(
              voc_size, dim, num_indices, dtype=dtype, device=device
          )
          torch.set_rng_state(rng_state)
          param1.grad = gen_sparse_grad(
              voc_size, dim, num_indices, dtype=dtype, device=device
          )
          optimizer.step()
          optimizer1.step()
  
      check_equal(param, param1, atol=1e-2, verbose=True)
  
file_path: DeepSpeed/tests/unit/ops/adagrad/test_cpu_adagrad.py
imports: import torch
  import numpy
  import pytest
  import deepspeed
  from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import CPUAdagradBuilder
  from unit.common import DistributedTest
label_desc: Use image node_id 2 for calling the TestCPUAdagrad obj's underlying member method code with example usage: obj.test_cpu_adagrad_opt_sparse_embedding(model_size, vocabulary_size, dim) and returns: new_t