3.1.21.1.3. unit_scaling.core.functional.scale_elementwise

unit_scaling.core.functional.scale_elementwise(f: Callable[[...], Tensor], output_scale: float, grad_input_scale: float, constraint: str | None = 'to_output_scale') → Callable[[...], Tensor]

Transforms an element-wise function into a scaled version.

Parameters:

• f (Callable[..., Tensor]) – the element-wise function to be scaled. Should take as its first input a Tensor, followed by *args, **kwargs.

• output_scale (float) – the scale to be applied to the output

• grad_input_scale (float) – the scale to be applied to the grad of the input

• constraint (Optional[str], optional) – The name of the constraint function to be applied to the outputs & input gradient. In this case, the constraint name must be one of: [None, ‘gmean’, ‘hmean’, ‘amean’, ‘to_output_scale’, ‘to_grad_input_scale’] (see unit_scaling.constraints for details on these constraint functions). Defaults to gmean.

Returns:

the scaled function

Return type:

Callable[…, Tensor]