nsl.lib.project_to_ball

Projects batched tensors to a ball with the given radius in the given norm.

nsl.lib.project_to_ball(
    tensors, radius, norm_type, epsilon=1e-06
)

For each example in the batch, its global norm (across all tensors) will be clipped to the given raidus. If its global norm is already smaller or equal to the given radius, its values won't be changed.

Only L-infinity and L2 norms are currently supported.

tensors A (nested) collection of batched tensors to project to the given ball. The first dimension of each tensor is the batch size and must all be equal.
radius the radius of the ball.
norm_type One of nsl.configs.NormType. Currently L1 norm is not supported.
epsilon Used to avoid division by 0.

A collection of tensors in the same structure as the input, projected to the given ball.