Add encoding type preservation support for float16 and bfloat16.

PiperOrigin-RevId: 515742571

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/stages_impl.py

@@ -475,10 +475,10 @@ def encode(self, x, encode_params):
     # type to be able to recover the type from encoded_tensors in decode method.
     if x.dtype == tf.float32:
       return collections.OrderedDict([(self.ENCODED_VALUES_KEY, packed_x)])
-    elif x.dtype == tf.float64:
+    elif x.dtype in [tf.float16, tf.bfloat16, tf.float64]:
       return collections.OrderedDict([(self.ENCODED_VALUES_KEY, packed_x),
                                       (self.DUMMY_TYPE_VALUES_KEY,
-                                       tf.constant(0.0, dtype=tf.float64))])
+                                       tf.constant(0.0, dtype=x.dtype))])
     else:
       raise TypeError(
           'Unsupported packing type: %s. Supported types are tf.float32 and '

tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/stages_impl_test.py

@@ -102,21 +102,23 @@ def is_lossless(self):
     """See base class."""
     return True
 
-  def common_asserts_for_test_data(self, data):
+  def common_asserts_for_test_data(self, data, rtol=1e-6, atol=1e-6):
     """See base class."""
     encoded_x = data.encoded_x[
         stages_impl.HadamardEncodingStage.ENCODED_VALUES_KEY]
-    self.assertAllClose(data.x, data.decoded_x)
+    self.assertAllClose(data.x, data.decoded_x, rtol=rtol, atol=atol)
     self.assertLen(encoded_x.shape, 2)
     # This is a rotation, hence, the norms should be the same.
     # If the input has dimension 1, the transform is applied to the whole input.
     # If the input has dimension 2, the transform is applied to every single
     # vector separately.
     if len(data.x.shape) == 1:
-      self.assertAllClose(np.linalg.norm(data.x), np.linalg.norm(encoded_x))
+      self.assertAllClose(np.linalg.norm(data.x), np.linalg.norm(encoded_x),
+                          rtol=rtol, atol=atol)
     else:
       for x, y in zip(data.x, encoded_x):
-        self.assertAllClose(np.linalg.norm(x), np.linalg.norm(y))
+        self.assertAllClose(np.linalg.norm(x), np.linalg.norm(y),
+                            rtol=rtol, atol=atol)
 
   def test_encoding_randomized(self):
     # The encoding stage declares a source of randomness (a random seed) in the
@@ -174,12 +176,15 @@ def get_random_shape_input():
     self.assertEqual(test_data.x.shape[0], encoded_shape[0])
     self.assertEqual(8, encoded_shape[1])
 
-  @parameterized.parameters([tf.float32, tf.float64])
-  def test_input_types(self, x_dtype):
+  @parameterized.parameters([(tf.float16, 1e-3, 1e-3),
+                             (tf.bfloat16, 1e-1, 1e-1),
+                             (tf.float32, 1e-6, 1e-6),
+                             (tf.float64, 1e-6, 1e-6)])
+  def test_input_types(self, x_dtype, rtol, atol):
     test_data = self.run_one_to_many_encode_decode(
         self.default_encoding_stage(),
         lambda: tf.random.normal([1, 12], dtype=x_dtype))
-    self.common_asserts_for_test_data(test_data)
+    self.common_asserts_for_test_data(test_data, rtol=rtol, atol=atol)
 
   def test_unknown_shape_raises(self):
     x = test_utils.get_tensor_with_random_shape()