code: def _batch_form_diag(tsr):
      """Form diagonal matrices in batch mode."""
      eye = torch.eye(tsr.shape[-1], dtype=tsr.dtype, device=tsr.device)
      M = tsr.unsqueeze(-1).expand(tsr.shape + tsr.shape[-1:])
      return eye * M
  
imports: import torch
  from linear_operator.utils.cholesky import psd_safe_cholesky
  from torch.distributions import constraints, LKJCholesky
  from torch.nn import Module
  from ..None import settings
  from .prior import Prior
label_desc: Use image node_id 1 for calling a global function with example usage: _batch_form_diag(tsr) and returns: unknown