3.1.22.9. unit_scaling.functional.linear_readout

unit_scaling.functional.linear_readout(input: Tensor, weight: Tensor, bias: Tensor | None, constraint: str | None = None) → Tensor[source]

Applies a unit-scaled linear transformation, for the final network output.

Applies a linear transformation to the incoming data: \(y = xA^T + b\).

This operation supports 2-D weight with sparse layout

Warning

Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, or may not have autograd support. If you notice missing functionality please open a feature request.

This operator supports TensorFloat32.

Parameters:

constraint (Optional[str]?) – 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.

Shape:

  • Input: \((*, in\_features)\) where * means any number of additional dimensions, including none

  • Weight: \((out\_features, in\_features)\) or \((in\_features)\)

  • Bias: \((out\_features)\) or \(()\)

  • Output: \((*, out\_features)\) or \((*)\), based on the shape of the weight