UnitNormalization

Normalize a batch of inputs so that each input in the batch has a L2 norm equal to 1.

Abstract Signature:

UnitNormalization(axis: int = -1)

PyTorchKerasTensorFlowFlax NNX

PyTorch

API: —

Strategy: Plugin (unit_norm_layer)

Keras

API: keras.layers.UnitNormalization

Strategy: Direct Mapping

Official Docs ↗

TensorFlow

API: tf.keras.layers.UnitNormalization

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: —

Strategy: Plugin (unit_norm_layer)