code: def _rank3_diag(x: Tensor) -> Tensor:
      eye = torch.eye(x.size(1)).type_as(x)
      out = eye * x.unsqueeze(2).expand(x.size(0), x.size(1), x.size(1))
  
      return out
  
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
label_desc: Use image node_id 3 for calling a global function with example usage: _rank3_diag(x) and returns: out