Adds as_gather_encoder utility.

PiperOrigin-RevId: 264910823

Author: PraChetit

tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders/BUILD

@@ -19,6 +19,7 @@ py_library(
     deps = [
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:core_encoder",
+        "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:gather_encoder",
         "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:simple_encoder",
         "//tensorflow_model_optimization/python/core/internal/tensor_encoding/stages:stages_impl",
     ],
@@ -34,6 +35,7 @@ py_test(
         # tensorflow dep1,
         # python:util tensorflow dep2,
         "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:core_encoder",
+        "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:gather_encoder",
         "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:simple_encoder",
     ],
 )

tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders/__init__.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import as_gather_encoder
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import as_simple_encoder
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import hadamard_quantization
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import identity

tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders/common_encoders.py

@@ -24,6 +24,7 @@
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import core_encoder
+from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import gather_encoder
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import simple_encoder
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages import stages_impl
 
@@ -33,8 +34,8 @@ def as_simple_encoder(encoder, tensorspec):
 
   Args:
     encoder: An `Encoder` object to be used to encoding.
-    tensorspec: A `TensorSpec`. The created `SimpleEncoderV2` will be
-      constrained to only encode input values compatible with `tensorspec`.
+    tensorspec: A `TensorSpec`. The created `SimpleEncoder` will be constrained
+      to only encode input values compatible with `tensorspec`.
 
   Returns:
     A `SimpleEncoder`.
@@ -50,6 +51,28 @@ def as_simple_encoder(encoder, tensorspec):
   return simple_encoder.SimpleEncoder(encoder, tensorspec)
 
 
+def as_gather_encoder(encoder, tensorspec):
+  """Wraps an `Encoder` object as a `GahterEncoder`.
+
+  Args:
+    encoder: An `Encoder` object to be used to encoding.
+    tensorspec: A `TensorSpec`. The created `GahterEncoder` will be constrained
+      to only encode input values compatible with `tensorspec`.
+
+  Returns:
+    A `GahterEncoder`.
+
+  Raises:
+    TypeError:
+      If `encoder` is not an `Encoder` or `tensorspec` is not a `TensorSpec`.
+  """
+  if not isinstance(encoder, core_encoder.Encoder):
+    raise TypeError('The encoder must be an instance of `Encoder`.')
+  if not isinstance(tensorspec, tf.TensorSpec):
+    raise TypeError('The tensorspec must be a tf.TensorSpec.')
+  return gather_encoder.GatherEncoder.from_encoder(encoder, tensorspec)
+
+
 def identity():
   """Returns identity `Encoder`."""
   return core_encoder.EncoderComposer(

tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders/common_encoders_test.py

@@ -21,6 +21,7 @@
 
 from tensorflow.python.util import nest as core_nest
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import core_encoder
+from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import gather_encoder
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import simple_encoder
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders import common_encoders
 
@@ -52,6 +53,24 @@ def test_as_simple_encoder_raises_tensorspec(self, not_a_tensorspec):
       common_encoders.as_simple_encoder(common_encoders.identity(),
                                         not_a_tensorspec)
 
+  @parameterized.parameters(_ENCODER_FNS)
+  def test_as_gather_encoder(self, encoder_fn):
+    encoder = common_encoders.as_gather_encoder(encoder_fn(),
+                                                tf.TensorSpec((2,), tf.float32))
+    self.assertIsInstance(encoder, gather_encoder.GatherEncoder)
+
+  @parameterized.parameters(None, [[]], 2.0, 'string')
+  def test_as_gather_encoder_raises_encoder(self, not_an_encoder):
+    with self.assertRaises(TypeError):
+      common_encoders.as_gather_encoder(not_an_encoder,
+                                        tf.TensorSpec((2,), tf.float32))
+
+  @parameterized.parameters(None, [[]], 2.0, 'string')
+  def test_as_gather_encoder_raises_tensorspec(self, not_a_tensorspec):
+    with self.assertRaises(TypeError):
+      common_encoders.as_gather_encoder(common_encoders.identity(),
+                                        not_a_tensorspec)
+
   def test_identity(self):
     encoder = common_encoders.identity()
     self.assertIsInstance(encoder, core_encoder.Encoder)