docstr: Calculates the trace of a matrix and is able to do broadcasting over batch
  dimensions, unlike `torch.trace`.
  
  Broadcasting is necessary for the conditional version of the transform,
  where `self.weights` may have batch dimensions corresponding the batch
  dimensions of the context variable that was conditioned upon.
docstr_tok: ["Calculates","the","trace","of","a","matrix","and","is","able","to","do","broadcasting","over","batch","dimensions",",","unlike","`","torch.trace","`",".","Broadcasting","is","necessary","for","the","conditional","version","of","the","transform",",","where","`","self.weights","`","may","have","batch","dimensions","corresponding","the","batch","dimensions","of","the","context","variable","that","was","conditioned","upon","."]
code: def _trace(self, M):
      """
      Calculates the trace of a matrix and is able to do broadcasting over batch
      dimensions, unlike `torch.trace`.
  
      Broadcasting is necessary for the conditional version of the transform,
      where `self.weights` may have batch dimensions corresponding the batch
      dimensions of the context variable that was conditioned upon.
      """
      return M.diagonal(dim1=-2, dim2=-1).sum(-1)
  
imports: import math
  from functools import partial
  import torch
  import torch.nn
  from torch.distributions import Transform, constraints
  from pyro.nn import DenseNN
  from ..conditional import ConditionalTransformModule
  from ..torch_transform import TransformModule
  from ..util import copy_docs_from
label_desc: Use image node_id 3 for calling the ConditionedMatrixExponential obj's underlying member method code with example usage: obj._trace(M) and returns: M