code: def test_clip_grad_norm(max_norm, check_isfinite):
      def gt_grad_global_norm(parameters):
          ret = 0
          for p in parameters:
              grads = p.list_grad()
              ret += (grads[0].asnumpy() ** 2).sum()
          return np.sqrt(ret)
  
      contexts = [mx.cpu(0), mx.cpu(1)]
      net = mx.gluon.nn.HybridSequential()
      # Create a network with 8 layers
      for _ in range(8):
          net.add(
              mx.gluon.nn.Dense(
                  1, weight_initializer="ones", bias_initializer="ones"
              )
          )
      net.initialize(ctx=contexts)
      net.hybridize()
      trainer = mx.gluon.Trainer(
          net.collect_params(), "sgd", update_on_kvstore=False
      )
      for ctx in contexts:
          with mx.autograd.record():
              out = net(mx.np.ones((1, 1), ctx=ctx))
          out.backward()
      trainer.allreduce_grads()
      # Cache the original gradient for checking
      original_grad_l = [
          p.list_grad()[0].asnumpy()
          for p in net.collect_params().values()
      ]
      norm = grad_global_norm(net.collect_params().values())
      gt_norm = gt_grad_global_norm(net.collect_params().values())
      assert_almost_equal(norm, gt_norm, atol=1e-5)
      with mx.cpu(2):
          norm, ratio, is_finite = clip_grad_global_norm(
              net.collect_params().values(), max_norm, check_isfinite
          )
      assert_almost_equal(norm, gt_norm, atol=1e-5)
      for p, orig_grad in zip(
          net.collect_params().values(), original_grad_l
      ):
          for ctx in contexts:
              if max_norm > norm:
                  assert_almost_equal(p.grad(ctx).asnumpy(), orig_grad)
              else:
                  ratio = max_norm / norm
                  assert_almost_equal(
                      p.grad(ctx).asnumpy(), orig_grad * ratio
                  )
  
file_path: gluon-nlp/tests/test_utils_parameter.py
imports: import pytest
  import mxnet
  import numpy
  from numpy.testing import assert_allclose
  from mxnet.gluon import nn
  from gluonnlp.utils.parameter import grad_global_norm, clip_grad_global_norm, AverageSGDTracker
  from mxnet.test_utils import assert_almost_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_clip_grad_norm(max_norm, check_isfinite) and returns: np