Skip to content

Math¤

distreqx.utils.math.multiply_no_nan(x: Array, y: Array) -> Array ¤

Computes the element-wise product of x and y, returning 0 where y is zero, even if x is NaN or infinite.

Arguments:

  • x: First input.
  • y: Second input.

Returns:

  • The product of x and y.

Raises:

  • ValueError if the shapes of x and y do not match.

distreqx.utils.math.power_no_nan(x: Array, y: Array) -> Array ¤

Computes \(x^y\), ensuring the result is \(1\) when \(y = 0\), following the convention \(0^0 = 1\).

Arguments:

  • x: First input.
  • y: Second input.

Returns:

  • The power \(x^y\).

distreqx.utils.math.mul_exp(x: Array, logp: Array) -> Array ¤

Returns \(x\exp(\ell)\) with zero output if \(\exp(\ell) = 0\).

Arguments:

  • x: An array.
  • logp: An array representing logarithms, denoted \(\ell\).

Returns:

  • The result \(x\exp(\ell)\).

distreqx.utils.math.normalize(*, probs: Array | None = None, logits: Array | None = None) -> Array ¤

Normalizes logits via log_softmax or probabilities to ensure they sum to one.

Arguments:

  • probs: Probability values.
  • logits: Logit values.

Returns:

  • Normalized probabilities or logits.

distreqx.utils.math.sum_last(x: Array, ndims: int) -> Array ¤

Sums the last ndims axes of array x.

Arguments:

  • x: An array.
  • ndims: The number of last dimensions to sum.

Returns:

  • The sum of the last ndims dimensions of x.

distreqx.utils.math.log_expbig_minus_expsmall(big: Array, small: Array) -> Array ¤

Stable implementation of \(\log(e^b - e^s)\).

Arguments:

  • big: First input, denoted \(b\).
  • small: Second input, denoted \(s\). It must satisfy \(s \le b\).

Returns:

  • The resulting \(\log(e^b - e^s)\).

distreqx.utils.math.log_beta(a: Array, b: Array) -> Array ¤

Obtains the log of the beta function \(\log B(a, b)\).

Arguments:

  • a: First input. It must be positive.
  • b: Second input. It must be positive.

Returns:

The value

\[ \log B(a, b) = \log\Gamma(a) + \log\Gamma(b) - \log\Gamma(a + b) \]

where \(\Gamma\) is the Gamma function, obtained through stable computation of \(\log\Gamma\).

distreqx.utils.math.log_beta_multivariate(a: Array) -> Array ¤

Obtains the log of the multivariate beta function \(\log B(a)\).

Arguments:

  • a: An array of length \(K\) containing positive values.

Returns:

The value

\[ \log B(a) = \sum_{k=1}^{K} \log\Gamma(a_k) - \log\Gamma\left(\sum_{k=1}^{K} a_k\right) \]

where \(\Gamma\) is the Gamma function, obtained through stable computation of \(\log\Gamma\).