code: def _softmax(layout, values):
      h, r, c = layout.nonzero(as_tuple=True)
      norms = torch.logsumexp(values, dim=-1, keepdim=True)
      linear_idx = h * layout.shape[1] + r
  
      out_t = torch.zeros(
          norms.shape[0],
          layout.shape[0] * layout.shape[1],
          norms.shape[2],
          norms.shape[3],
          dtype=norms.dtype,
          device=norms.device,
      )
      max_val = norms.max()
      out_t.index_add_(
          1, linear_idx.to(values.device), (norms - max_val).exp()
      ).clamp_min_(1e-24).log_().add_(max_val)
      out = torch.exp(values - out_t[:, linear_idx])
      return out
  
imports: import logging
  import torch
  from xformers import _is_triton_available
  from xformers.ops import masked_matmul
label_desc: Use image node_id 3 for calling a global function with example usage: _softmax(layout, values) and returns: out