Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/layers/normalization/unit_normalization.py: 46%
26 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Unit Normalization layer."""
18import tensorflow.compat.v2 as tf
20from keras.src.engine import base_layer
21from keras.src.utils import tf_utils
23# isort: off
24from tensorflow.python.util.tf_export import keras_export
27@keras_export("keras.layers.UnitNormalization", v1=[])
28class UnitNormalization(base_layer.Layer):
29 """Unit normalization layer.
31 Normalize a batch of inputs so that each input in the batch has a L2 norm
32 equal to 1 (across the axes specified in `axis`).
34 Example:
36 >>> data = tf.constant(np.arange(6).reshape(2, 3), dtype=tf.float32)
37 >>> normalized_data = tf.keras.layers.UnitNormalization()(data)
38 >>> print(tf.reduce_sum(normalized_data[0, :] ** 2).numpy())
39 1.0
41 Args:
42 axis: Integer or list/tuple. The axis or axes to normalize across.
43 Typically this is the features axis or axes. The left-out axes are
44 typically the batch axis or axes. Defaults to `-1`, the last dimension
45 in the input.
46 """
48 def __init__(self, axis=-1, **kwargs):
49 super().__init__(**kwargs)
50 if isinstance(axis, (list, tuple)):
51 self.axis = list(axis)
52 elif isinstance(axis, int):
53 self.axis = axis
54 else:
55 raise TypeError(
56 "Invalid value for `axis` argument: "
57 "expected an int or a list/tuple of ints. "
58 f"Received: axis={axis}"
59 )
60 self.supports_masking = True
62 def build(self, input_shape):
63 self.axis = tf_utils.validate_axis(self.axis, input_shape)
65 def call(self, inputs):
66 inputs = tf.cast(inputs, self.compute_dtype)
67 return tf.linalg.l2_normalize(inputs, axis=self.axis)
69 def compute_output_shape(self, input_shape):
70 return input_shape
72 def get_config(self):
73 config = super().get_config()
74 config.update({"axis": self.axis})
75 return config