def METHOD_NAME(g, idtype):
    g = g.astype(idtype).to(F.ctx())
    g.ndata["h"] = F.randn((g.num_nodes(), 3))
    g.edata["h"] = F.randn((g.num_edges(), 2))
    # Test.1: node readout
    x = dgl.softmax_nodes(g, "h")
    subg = dgl.unbatch(g)
    subx = []
    for sg in subg:
        subx.append(F.softmax(sg.ndata["h"], dim=0))
    assert F.allclose(x, F.cat(subx, dim=0))
    # Test.2: edge readout
    x = dgl.softmax_edges(g, "h")
    subg = dgl.unbatch(g)
    subx = []
    for sg in subg:
        subx.append(F.softmax(sg.edata["h"], dim=0))
    assert F.allclose(x, F.cat(subx, dim=0))