clamped_length_sqr#

iskra.geometry.distances.clamped_length_sqr(x: Tensor, dim: int | tuple[int, ...] = -1, keepdim: bool = False, eps: float = 1e-12) Tensor[SOURCE]#

Computes the _squared_ lengths of a set of vectors and then clamps them.

Helps avoid numerical issues when computing lengths of vectors:

  • The gradient of the square-root tends to infinitey as we get closer to zero.

  • We often wish to divide by squared length, which also explodes around zero.

Parameters:
  • x (Tensor[Float, [Bs, Dim]]) – Set of vectors.

  • dim (int | tuple[int, ...]) – The dimension(s) along which to compute lengths.

  • keepdim (bool) – Whether to reduce the selected dimensions or to keep them with the length one.

  • eps (float) – The _squared length_ will be clamped to this minimum value.

Returns:

(Tensor[Float, [Bs] | [Bs, 1]]) – Squared lengths of vectors along dimension dim. Last dimension of output is 1 if keepdim=True, otherwise it is removed.