Normalization Functionals#

physicsnemo.nn.functional.safe_normalize(
vectors: Float[Tensor, '...'],
dim: int,
*,
implementation: Literal['torch'] | None = None,
) → Float[Tensor, '...']#

Scale vectors to unit L2 length along dim, preserving zero vectors.

Each vector is divided by its largest absolute component before computing its norm. This avoids overflow and underflow from the input magnitude without an absolute epsilon floor that would shorten small vectors.

Parameters:

  • vectors (Float[torch.Tensor, "..."]) – Floating-point vectors of any shape.

  • dim (int) – Dimension holding the vector components.

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

Returns:

Unit vectors with the input shape, device, and dtype, including under autocast. Exactly zero vectors remain zero.

Return type:

Float[torch.Tensor, “…”]

Notes

Non-finite components propagate NaNs to the whole vector. Derivatives near zero can exceed the dtype’s range even when forward values are finite.

Examples

>>> v = torch.tensor([[3.0e-13, 4.0e-13], [0.0, 0.0]])
>>> safe_normalize(v, dim=-1)
tensor([[0.6000, 0.8000],
        [0.0000, 0.0000]])
On this page