code: def trace(x: Tensor) -> Tensor:
      """Compute trace of input tensor (batched)"""
      b, m, n = x.size()
      assert m == n
      mask = (
          torch.eye(n, dtype=torch.bool, device=x.device)
          .unsqueeze(0)
          .expand_as(x)
      )
      trace = (
          x.masked_select(mask)
          .contiguous()
          .view(b, n)
          .sum(dim=-1, keepdim=False)
      )
      return trace
  
imports: import torch
  from torch import Tensor
  from torch.nn import functional
label_desc: Use image node_id 2 for calling a global function with example usage: trace(x) and returns: trace