code: def test_trainer_sparse_kv():
      def check_trainer_sparse_kv(
          kv, stype, grad_stype, update_on_kv, expected
      ):
          x = mx.gluon.Parameter(
              "x",
              shape=(10, 1),
              lr_mult=1.0,
              stype=stype,
              grad_stype=grad_stype,
          )
          x.initialize(ctx=[mx.cpu(0), mx.cpu(1)], init="zeros")
          trainer = gluon.Trainer(
              [x],
              "sgd",
              {"learning_rate": 0.1},
              kvstore=kv,
              update_on_kvstore=update_on_kv,
          )
          all_rows = mx.nd.arange(0, 10, ctx=mx.cpu(0))
          try:
              ws = (
                  x.list_data()
                  if stype == "default"
                  else x.list_row_sparse_data(all_rows)
              )
              with mx.autograd.record():
                  for w in ws:
                      y = w + 1
                      y.backward()
              trainer.step(1)
              assert trainer._kvstore.type == kv
              assert trainer._kv_initialized
              assert trainer._update_on_kvstore is expected
              # the updated parameter should be based on the loaded checkpoint
              mx.nd.waitall()
              updated_w = (
                  x.data(mx.cpu(0))
                  if stype == "default"
                  else x.row_sparse_data(all_rows)
              )
              assert (updated_w == -0.2).asnumpy().all(), updated_w
          except Exception as err:
              assert isinstance(err, expected)
  
      kvs = ["local", "device"]
      global_update_on_kvstore = bool(
          int(os.getenv("MXNET_UPDATE_ON_KVSTORE", "1"))
      )
      for kv in kvs:
          check_trainer_sparse_kv(kv, "default", "default", True, True)
          check_trainer_sparse_kv(
              kv, "default", "default", False, False
          )
          check_trainer_sparse_kv(
              kv, "default", "default", None, global_update_on_kvstore
          )
          check_trainer_sparse_kv(
              kv, "default", "row_sparse", None, False
          )
          check_trainer_sparse_kv(
              kv, "default", "row_sparse", True, True
          )
          check_trainer_sparse_kv(
              kv, "default", "row_sparse", False, False
          )
          check_trainer_sparse_kv(
              kv, "row_sparse", "row_sparse", None, True
          )
          check_trainer_sparse_kv(
              kv, "row_sparse", "row_sparse", False, ValueError
          )
  
imports: import mxnet
  import unittest
  import os
  import numpy
  from mxnet import gluon
  from mxnet.gluon import nn
  from mxnet.test_utils import assert_almost_equal
  from common import assertRaises, xfail_when_nonstandard_decimal_separator
  from copy import deepcopy
  import pytest
label_desc: Use image node_id 9 for calling a global function with example usage: test_trainer_sparse_kv() without return types