code: def test_correctness(self):
      layer = layers.UnitNormalization(axis=-1)
      inputs = np.random.normal(size=(2, 3))
      outputs = layer(inputs)
      self.assertAllClose(squared_l2_norm(outputs[0, :]), 1.0)
      self.assertAllClose(squared_l2_norm(outputs[1, :]), 1.0)
  
      layer = layers.UnitNormalization(axis=(1, 2))
      inputs = np.random.normal(size=(2, 3, 3))
      outputs = layer(inputs)
      self.assertAllClose(squared_l2_norm(outputs[0, :, :]), 1.0)
      self.assertAllClose(squared_l2_norm(outputs[1, :, :]), 1.0)
  
      layer = layers.UnitNormalization(axis=1)
      inputs = np.random.normal(size=(2, 3, 2))
      outputs = layer(inputs)
      self.assertAllClose(squared_l2_norm(outputs[0, :, 0]), 1.0)
      self.assertAllClose(squared_l2_norm(outputs[1, :, 0]), 1.0)
      self.assertAllClose(squared_l2_norm(outputs[0, :, 1]), 1.0)
      self.assertAllClose(squared_l2_norm(outputs[1, :, 1]), 1.0)
  
file_path: keras/keras/layers/normalization/unit_normalization_test.py
imports: import numpy
  import pytest
  from keras import backend
  from keras import layers
  from keras import testing
label_desc: Use image node_id 3 for calling the UnitNormalizationTest obj's underlying member method code with example usage: obj.test_correctness() without return types