code: def test_arbitrary_dimension(dim):
      shape = [3, 4, 2, 5]
      X = torch.randn(*shape, dtype=torch.float64)
  
      alpha_shape = shape
      alpha_shape[dim] = 1
  
      alphas = 0.05 + 2.5 * torch.rand(alpha_shape, dtype=torch.float64)
  
      P = entmax_bisect(X, alpha=alphas, dim=dim)
  
      ranges = [
          list(range(k)) if i != dim else [slice(None)]
          for i, k in enumerate(shape)
      ]
  
      for ix in product(*ranges):
          x = X[ix].unsqueeze(0)
          alpha = alphas[ix].item()
          p_true = entmax_bisect(x, alpha=alpha, dim=-1)
          assert torch.allclose(P[ix], p_true)
  
file_path: ludwig/tests/ludwig/utils/entmax/test_root_finding.py
imports: from functools import partial
  from itertools import product
  import pytest
  import torch
  from torch.autograd import gradcheck
  from ludwig.utils.entmax.activations import entmax15, sparsemax
  from ludwig.utils.entmax.root_finding import entmax_bisect, sparsemax_bisect
label_desc: Use image node_id 7 for calling a global function with example usage: test_arbitrary_dimension(dim) without return types