code: def test_edge_softmax(g, norm_by, shp, idtype):
      g = g.astype(idtype).to(F.ctx())
      edata = F.tensor(np.random.rand(g.num_edges(), *shp))
      e1 = F.attach_grad(F.clone(edata))
  
      with F.record_grad():
          score1 = edge_softmax(g, e1, norm_by=norm_by)
          F.backward(F.reduce_sum(score1))
          grad_edata = F.grad(e1)
  
      with F.record_grad():
          e2 = F.attach_grad(F.clone(edata))
          e2_2d = F.reshape(
              e2,
              (
                  g.number_of_src_nodes(),
                  g.number_of_dst_nodes(),
                  *e2.shape[1:],
              ),
          )
          if norm_by == "src":
              score2 = F.softmax(e2_2d, 1)
              score2 = F.reshape(score2, (-1, *e2.shape[1:]))
          if norm_by == "dst":
              score2 = F.softmax(e2_2d, 0)
              score2 = F.reshape(score2, (-1, *e2.shape[1:]))
          assert F.allclose(score1, score2)
          print("forward passed")
  
          F.backward(F.reduce_sum(score2))
          assert F.allclose(F.grad(e2), grad_edata)
          print("backward passed")
  
file_path: dgl/tests/python/common/ops/test_edge_softmax.py
imports: import itertools
  import math
  import unittest
  from collections import Counter
  import backend
  import dgl
  import dgl.function
  import networkx
  import numpy
  import pytest
  import scipy.sparse
  from dgl import DGLError
  from dgl.ops import edge_softmax
  from scipy.sparse import rand
  from utils import get_cases, parametrize_idtype
label_desc: Use image node_id 1 for calling a global function with example usage: test_edge_softmax(g, norm_by, shp, idtype) without return types