code: def stable_softmax(t, dim=-1, alpha=32**2):
      t = t / alpha
      t = t - torch.amax(t, dim=dim, keepdim=True).detach()
      return (t * alpha).softmax(dim=dim)
  
imports: import torch
  import torch.nn.functional
  from einops import rearrange
  from torch import nn
label_desc: Use image node_id 5 for calling a global function with example usage: stable_softmax(t, dim, alpha) and returns: unknown