code: def test_softmax_backward():
      src_sparse = torch.rand(4, 8)
      index = torch.tensor([0, 0, 1, 1])
      src_dense = src_sparse.clone().view(2, 2, src_sparse.size(-1))
  
      src_sparse.requires_grad_(True)
      src_dense.requires_grad_(True)
  
      out_sparse = softmax(src_sparse, index)
      out_sparse.mean().backward()
      out_dense = src_dense.softmax(dim=1)
      out_dense.mean().backward()
  
      assert torch.allclose(out_sparse, out_dense.view_as(out_sparse))
      assert torch.allclose(
          src_sparse.grad, src_dense.grad.view_as(src_sparse)
      )
  
imports: import pytest
  import torch
  import torch_geometric.typing
  from torch_geometric.profile import benchmark
  from torch_geometric.utils import softmax
label_desc: Use image node_id 2 for calling a global function with example usage: test_softmax_backward() without return types