code: def mode_of_distribution(distrib):
      assert isinstance(distrib, dists.Distribution)
      if isinstance(distrib, dists.Categorical):
          return distrib.probs.argmax(dim=-1)
      elif isinstance(
          distrib, (dists.Normal, dists.MultivariateNormal)
      ):
          return distrib.mean
      elif isinstance(
          distrib,
          dists.transformed_distribution.TransformedDistribution,
      ):
          x = mode_of_distribution(distrib.base_dist)
          for transform in distrib.transforms:
              x = transform(x)
          return x
      elif isinstance(distrib, torch.distributions.Independent):
          return mode_of_distribution(distrib.base_dist)
      else:
          raise RuntimeError("{} is not supported".format(distrib))
  
label_desc: Use image node_id 1 for calling a global function with example usage: mode_of_distribution(distrib) and returns: distrib, distrib, x, mode_of_distribution