docstr: Generate a random matrix whose rows are exactly orthonormal
  
  "How to generate random matrices from the classical compact groups", Mezzadri, 2007
  https://arxiv.org/pdf/math-ph/0609050v2.pdf
  
  .. note: the typical qr decomposition does not give uniform results, qr decomposition is not
  unique and the qr decomposition routines are biased towards numerical stability. See the above
  paper for more information.
  
  .. note: this does not follow the original implementation from the Performers authors.
  see docs/assets/kde plots to visualize the impact of using the R signs to correct Q
docstr_tok: ["Generate","a","random","matrix","whose","rows","are","exactly","orthonormal","''","How","to","generate","random","matrices","from","the","classical","compact","groups","''",",","Mezzadri",",","2007","https",":","\/\/arxiv.org\/pdf\/math-ph\/0609050v2.pdf","..","note",":","the","typical","qr","decomposition","does","not","give","uniform","results",",","qr","decomposition","is","not","unique","and","the","qr","decomposition","routines","are","biased","towards","numerical","stability",".","See","the","above","paper","for","more","information",".","..","note",":","this","does","not","follow","the","original","implementation","from","the","Performers","authors",".","see","docs\/assets\/kde","plots","to","visualize","the","impact","of","using","the","R","signs","to","correct","Q"]
code: def _get_random_ortho_matrix(
      blocks: int,
      dim: int,
      device: torch.device,
      norm_distribution: NormDistribution = NormDistribution.Uniform,
  ) -> torch.Tensor:
      r"""
      Generate a random matrix whose rows are exactly orthonormal
  
      "How to generate random matrices from the classical compact groups", Mezzadri, 2007
      https://arxiv.org/pdf/math-ph/0609050v2.pdf
  
      .. note: the typical qr decomposition does not give uniform results, qr decomposition is not
      unique and the qr decomposition routines are biased towards numerical stability. See the above
      paper for more information.
  
      .. note: this does not follow the original implementation from the Performers authors.
      see docs/assets/kde plots to visualize the impact of using the R signs to correct Q
      """
  
      H = torch.randn(
          (blocks, dim, dim), device=device, requires_grad=False
      )
  
      # Randomly scale the norms of the features, Xi distributed
      if norm_distribution == NormDistribution.Xi:
          # NOTE: This averages to sqrt(d)
          norms = torch.sqrt(torch.einsum("...d,...d->...", H, H))
  
      Q, R = torch.linalg.qr(H)
      Q = (
          torch.diag_embed(
              torch.sign(torch.diagonal(R, dim1=1, dim2=2))
          )
          @ Q
      )
  
      # Normalize if need be. Uniform NormDistribution does nothing, Q is already orthonormal
      if norm_distribution == NormDistribution.Xi:
          return torch.diag_embed(norms) @ Q
  
      return Q
  
imports: import math
  from enum import Enum, auto
  from typing import Optional
  import torch
  from torch.autograd.profiler import record_function
  from .base import FeatureMap
label_desc: Use image node_id 3 for calling the SoftMaxPositiveEstimators obj's underlying member method code with example usage: obj._get_random_ortho_matrix(blocks, dim, device, norm_distribution) and returns: Q, unknown