code: def orthogonal(
      module: Module,
      name: str = "weight",
      orthogonal_map: Optional[str] = None,
      *,
      use_trivialization: bool = True,
  ) -> Module:
      r"""Apply an orthogonal or unitary parametrization to a matrix or a batch of matrices.
  
      Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, the parametrized
      matrix :math:`Q \in \mathbb{K}^{m \times n}` is **orthogonal** as
  
      .. math::
  
          \begin{align*}
              Q^{\text{H}}Q &= \mathrm{I}_n \mathrlap{\qquad \text{if }m \geq n}\\
              QQ^{\text{H}} &= \mathrm{I}_m \mathrlap{\qquad \text{if }m < n}
          \end{align*}
  
      where :math:`Q^{\text{H}}` is the conjugate transpose when :math:`Q` is complex
      and the transpose when :math:`Q` is real-valued, and
      :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix.
      In plain words, :math:`Q` will have orthonormal columns whenever :math:`m \geq n`
      and orthonormal rows otherwise.
  
      If the tensor has more than two dimensions, we consider it as a batch of matrices of shape `(..., m, n)`.
  
      The matrix :math:`Q` may be parametrized via three different ``orthogonal_map`` in terms of the original tensor:
  
      - ``"matrix_exp"``/``"cayley"``:
        the :func:`~torch.matrix_exp` :math:`Q = \exp(A)` and the `Cayley map`_
        :math:`Q = (\mathrm{I}_n + A/2)(\mathrm{I}_n - A/2)^{-1}` are applied to a skew-symmetric
        :math:`A` to give an orthogonal matrix.
      - ``"householder"``: computes a product of Householder reflectors
        (:func:`~torch.linalg.householder_product`).
  
      ``"matrix_exp"``/``"cayley"`` often make the parametrized weight converge faster than
      ``"householder"``, but they are slower to compute for very thin or very wide matrices.
  
      If ``use_trivialization=True`` (default), the parametrization implements the "Dynamic Trivialization Framework",
      where an extra matrix :math:`B \in \mathbb{K}^{n \times n}` is stored under
      ``module.parametrizations.weight[0].base``. This helps the
      convergence of the parametrized layer at the expense of some extra memory use.
      See `Trivializations for Gradient-Based Optimization on Manifolds`_ .
  
      Initial value of :math:`Q`:
      If the original tensor is not parametrized and ``use_trivialization=True`` (default), the initial value
      of :math:`Q` is that of the original tensor if it is orthogonal (or unitary in the complex case)
      and it is orthogonalized via the QR decomposition otherwise (see :func:`torch.linalg.qr`).
      Same happens when it is not parametrized and ``orthogonal_map="householder"`` even when ``use_trivialization=False``.
      Otherwise, the initial value is the result of the composition of all the registered
      parametrizations applied to the original tensor.
  
      .. note::
          This function is implemented using the parametrization functionality
          in :func:`~torch.nn.utils.parametrize.register_parametrization`.
  
  
      .. _`Cayley map`: https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
      .. _`Trivializations for Gradient-Based Optimization on Manifolds`: https://arxiv.org/abs/1909.09501
  
      Args:
          module (nn.Module): module on which to register the parametrization.
          name (str, optional): name of the tensor to make orthogonal. Default: ``"weight"``.
          orthogonal_map (str, optional): One of the following: ``"matrix_exp"``, ``"cayley"``, ``"householder"``.
              Default: ``"matrix_exp"`` if the matrix is square or complex, ``"householder"`` otherwise.
          use_trivialization (bool, optional): whether to use the dynamic trivialization framework.
              Default: ``True``.
  
      Returns:
          The original module with an orthogonal parametrization registered to the specified
          weight
  
      Example::
  
          >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
          >>> orth_linear = orthogonal(nn.Linear(20, 40))
          >>> orth_linear
          ParametrizedLinear(
          in_features=20, out_features=40, bias=True
          (parametrizations): ModuleDict(
              (weight): ParametrizationList(
              (0): _Orthogonal()
              )
          )
          )
          >>> # xdoctest: +IGNORE_WANT
          >>> Q = orth_linear.weight
          >>> torch.dist(Q.T @ Q, torch.eye(20))
          tensor(4.9332e-07)
      """
      weight = getattr(module, name, None)
      if not isinstance(weight, Tensor):
          raise ValueError(
              f"Module '{module}' has no parameter or buffer with name '{name}'"
          )
  
      # We could implement this for 1-dim tensors as the maps on the sphere
      # but I believe it'd bite more people than it'd help
      if weight.ndim < 2:
          raise ValueError(
              "Expected a matrix or batch of matrices. "
              f"Got a tensor of {weight.ndim} dimensions."
          )
  
      if orthogonal_map is None:
          orthogonal_map = (
              "matrix_exp"
              if weight.size(-2) == weight.size(-1)
              or weight.is_complex()
              else "householder"
          )
  
      orth_enum = getattr(_OrthMaps, orthogonal_map, None)
      if orth_enum is None:
          raise ValueError(
              'orthogonal_map has to be one of "matrix_exp", "cayley", "householder". '
              f"Got: {orthogonal_map}"
          )
      orth = _Orthogonal(
          weight, orth_enum, use_trivialization=use_trivialization
      )
      parametrize.register_parametrization(
          module, name, orth, unsafe=True
      )
      return module
  
imports: from enum import Enum, auto
  import torch
  from torch import Tensor
  from ..utils import parametrize
  from ..modules import Module
  from ..None import functional
  from typing import Optional
label_desc: Use image node_id 3 for calling a global function with example usage: orthogonal(module, name, orthogonal_map) and returns: module