text: #
  # SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  # SPDX-License-Identifier: Apache-2.0
  #
  """Classes and functions related to MIMO transmit precoding"""
  
  import tensorflow as tf
  from sionna.utils import matrix_inv
  
  def zero_forcing_precoder(x, h, return_precoding_matrix=False):
      # pylint: disable=line-too-long
      r"""Zero-Forcing (ZF) Precoder
  
      This function implements ZF precoding for a MIMO link, assuming the
      following model:
  
      .. math::
  
          \mathbf{y} = \mathbf{H}\mathbf{G}\mathbf{x} + \mathbf{n}
  
      where :math:`\mathbf{y}\in\mathbb{C}^K` is the received signal vector,
      :math:`\mathbf{H}\in\mathbb{C}^{K\times M}` is the known channel matrix,
      :math:`\mathbf{G}\in\mathbb{C}^{M\times K}` is the precoding matrix,
      :math:`\mathbf{x}\in\mathbb{C}^K` is the symbol vector to be precoded,
      and :math:`\mathbf{n}\in\mathbb{C}^K` is a noise vector. It is assumed that
      :math:`K\le M`.
  
      The precoding matrix :math:`\mathbf{G}` is defined as (Eq. 4.37) [BHS2017]_ :
  
      .. math::
  
          \mathbf{G} = \mathbf{V}\mathbf{D}
  
      where
  
      .. math::
  
          \mathbf{V} &= \mathbf{H}^{\mathsf{H}}\left(\mathbf{H} \mathbf{H}^{\mathsf{H}}\right)^{-1}\\
          \mathbf{D} &= \mathop{\text{diag}}\left( \lVert \mathbf{v}_{k} \rVert_2^{-1}, k=0,\dots,K-1 \right).
  
      This ensures that each stream is precoded with a unit-norm vector,
      i.e., :math:`\mathop{\text{tr}}\left(\mathbf{G}\mathbf{G}^{\mathsf{H}}\right)=K`.
      The function returns the precoded vector :math:`\mathbf{G}\mathbf{x}`.
  
      Input
      -----
      x : [...,K], tf.complex
          1+D tensor containing the symbol vectors to be precoded.
  
      h : [...,K,M], tf.complex
          2+D tensor containing the channel matrices
  
      return_precoding_matrices : bool
          Indicates if the precoding matrices should be returned or not.
          Defaults to False.
  
      Output
      -------
      x_precoded : [...,M], tf.complex
          Tensor of the same shape and dtype as ``x`` apart from the last
          dimensions that has changed from `K` to `M`. It contains the
          precoded symbol vectors.
  
      g : [...,M,K], tf.complex
          2+D tensor containing the precoding matrices. It is only returned
          if ``return_precoding_matrices=True``.
  
      Note
      ----
      If you want to use this function in Graph mode with XLA, i.e., within
      a function that is decorated with ``@tf.function(jit_compile=True)``,
      you must set ``sionna.Config.xla_compat=true``.
      See :py:attr:`~sionna.Config.xla_compat`.
      """
  
      # Compute pseudo inverse for precoding
      g = tf.matmul(h, h, adjoint_b=True)
      g = tf.matmul(h, matrix_inv(g), adjoint_a=True)
  
      # Normalize each column to unit power
      norm = tf.sqrt(tf.reduce_sum(tf.abs(g)**2, axis=-2, keepdims=True))
      g = g/tf.cast(norm, g.dtype)
  
      # Expand last dim of `x` for precoding
      x_precoded = tf.expand_dims(x, -1)
  
      # Precode
      x_precoded = tf.squeeze(tf.matmul(g, x_precoded), -1)
  
      if return_precoding_matrix:
          return (x_precoded, g)
      else:
          return x_precoded
  
  
  
  