Changes fast_walsh_hadamard_transform to use tf.while_loop.

PiperOrigin-RevId: 263081103

Author: dengyinlin

tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/tf_utils.py

@@ -31,26 +31,52 @@ def fast_walsh_hadamard_transform(x):
 
   Args:
     x: A `Tensor`. Must be of shape `[a, b]`, where `a` can be anything (not
-      necessarily known), and `b` must be a power of two, statically known.
+      necessarily known), and `b` must be a power of two, not required to be
+      statically known.
 
   Returns:
     A `Tensor` of shape `[a, b]`, where `[i, :]` is the product `x[i, :]*H`,
       where `H` is the Hadamard matrix.
 
   Raises:
     ValueError: If the input is not rank 2 `Tensor`, and if the second dimension
-      is not a power of two.
+      is statically known and is not a power of two.
+    OpError: If the second dimension is not statically known and is not a power
+      of two. Note that in graph execution, this error is not raised during the
+      execution of the Python function, but during execution of the resulting
+      computation.
   """
   with tf.name_scope(None, 'fast_walsh_hadamard_transform'):
     # Validate input.
     x = tf.convert_to_tensor(x)
     if x.shape.ndims != 2:
       raise ValueError(
           'Number of dimensions of x must be 2. Shape of x: %s' % x.shape)
-    dim = x.shape.as_list()[1]
-    if not (dim and ((dim & (dim - 1)) == 0)):
-      raise ValueError('The dimension of x must be a power of two. '
-                       'Provided dimension is: %s' % dim)
+
+    original_x_shape = x.shape.as_list()
+    dim = x.shape.as_list()[-1]
+
+    if dim is None:  # dim is not statically known.
+      dim = tf.shape(x)[-1]
+      log2 = tf.cast(
+          tf.math.round(
+              tf.math.log(tf.cast(dim, tf.float32)) / tf.math.log(2.)),
+          tf.int32)
+      with tf.control_dependencies([
+          tf.compat.v1.assert_equal(
+              dim,
+              tf.math.pow(2, log2),
+              message='The dimension of x must be a power of two.'
+              'Provided dimension is: %s' % dim)
+      ]):
+        x = tf.identity(x)
+    else:  # dim is statically known.
+      if not (dim and ((dim & (dim - 1)) == 0)):
+        raise ValueError('The dimension of x must be a power of two. '
+                         'Provided dimension is: %s' % dim)
+      log2 = int(np.ceil(np.log2(dim)))
+      if dim == 1:  # Equivalent to identity.
+        return tf.identity(x)
 
     h_core = tf.constant([[1., 1.], [1., -1.]],
                          dtype=x.dtype,
@@ -60,17 +86,30 @@ def fast_walsh_hadamard_transform(x):
     # A step of the fast Walsh-Hadamard algorithm.
     def _hadamard_step(x, dim):
       """A single step in the fast Walsh-Hadamard transform."""
+      x_shape = x.shape.as_list()
       x = tf.reshape(x, [-1, 2])  # Reshape so that we have a matrix.
       x = tf.matmul(x, h_core)  # Multiply.
       x = tf.reshape(x, [-1, dim // 2, 2])  # Reshape to rank-3.
       x = tf.transpose(x, perm=permutation)  # Swap last two dimensions.
+      x.set_shape(x_shape)  # Failed shape inference in tf.while_loop.
       return x
 
-    # The fast Walsh-Hadamard transform.
-    for _ in range(int(np.ceil(np.log2(dim)))):
-      x = _hadamard_step(x, dim)
+    def _fwht(x, dim, log2):
+      x = tf.reshape(x, [-1, 2, dim // 2])
+      # The fast Walsh-Hadamard transform.
+
+      i = tf.constant(0)
+      c = lambda i, x: tf.less(i, log2)
+      b = lambda i, x: [i + 1, _hadamard_step(x, dim)]
+      i, x = tf.while_loop(c, b, [i, x])
+      return x
+
+    x = tf.cond(
+        tf.equal(dim, 1), lambda: tf.identity(x), lambda: _fwht(x, dim, log2))
+
     x = tf.reshape(x, [-1, dim])
     x /= tf.sqrt(tf.cast(dim, x.dtype))  # Normalize.
+    x.set_shape(original_x_shape)  # Failed shape inference after tf.while_loop.
     return x
 
 

tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/tf_utils_test.py

@@ -56,13 +56,61 @@ def test_illegal_inputs_shape(self, *dims):
       tf_utils.fast_walsh_hadamard_transform(x)
 
   @parameterized.parameters([[1, 3], [1, 7], [1, 9], [4, 3]])
-  def test_illegal_inputs_power_of_two(self, *dims):
-    """Tests incorrect shape of the rank 2 input."""
+  def test_illegal_inputs_static_power_of_two(self, *dims):
+    """Tests incorrect static shape of the rank 2 input."""
     x = tf.random.normal(dims)
     with self.assertRaisesRegexp(ValueError,
                                  'The dimension of x must be a power of two.'):
       tf_utils.fast_walsh_hadamard_transform(x)
 
+  def test_illegal_inputs_dynamic_power_of_two(self):
+    """Tests incorrect dynamic shape of the rank 2 input."""
+    rand = tf.random.uniform((), maxval=3, dtype=tf.int32)
+    x = tf.random.normal((3, 3**rand))
+    hx = tf_utils.fast_walsh_hadamard_transform(x)
+    with self.assertRaisesOpError('The dimension of x must be a power of two.'):
+      hx = self.evaluate(hx)
+
+  @parameterized.parameters([[1, 1], [4, 1], [2, 2], [1, 8], [1, 4]])
+  def test_static_input_shape(self, *dims):
+    """Tests static input shape."""
+    x = tf.random.normal(dims)
+    hx_tf = tf_utils.fast_walsh_hadamard_transform(x)
+    hhx_tf = tf_utils.fast_walsh_hadamard_transform(hx_tf)
+
+    x, hx_tf, hhx_tf = self.evaluate([x, hx_tf, hhx_tf])
+    self.assertAllEqual(x.shape, hhx_tf.shape)
+    self.assertAllClose(x, hhx_tf)
+
+  @parameterized.parameters([[1, 1], [4, 1], [2, 2], [1, 8], [1, 4]])
+  def test_static_input_output_shape(self, *dims):
+    """Tests static output shape is identical to static input shape."""
+    x = tf.random.normal(dims)
+    hx_tf = tf_utils.fast_walsh_hadamard_transform(x)
+    hhx_tf = tf_utils.fast_walsh_hadamard_transform(hx_tf)
+    self.assertEqual(list(dims), hx_tf.shape.as_list())
+    self.assertEqual(list(dims), hhx_tf.shape.as_list())
+
+  def test_dynamic_input_shape(self):
+    """Tests dynamic input shape."""
+    rand = tf.random.uniform((), maxval=4, dtype=tf.int32)
+    x = tf.random.normal((3, 2**rand))
+    hx_tf = tf_utils.fast_walsh_hadamard_transform(x)
+    hhx_tf = tf_utils.fast_walsh_hadamard_transform(hx_tf)
+    x, hx_tf, hhx_tf = self.evaluate([x, hx_tf, hhx_tf])
+    self.assertAllEqual(x.shape, hhx_tf.shape)
+    self.assertAllClose(x, hhx_tf)
+
+  def test_dynamic_input_shape_dim_one(self):
+    """Tests input shape where the second dimension is 1, dynamically known."""
+    rand = tf.random.uniform((), maxval=1, dtype=tf.int32)
+    x = tf.random.normal((3, 2**rand))
+    hx_tf = tf_utils.fast_walsh_hadamard_transform(x)
+    hhx_tf = tf_utils.fast_walsh_hadamard_transform(hx_tf)
+    x, hx_tf, hhx_tf = self.evaluate([x, hx_tf, hhx_tf])
+    self.assertAllEqual(x.shape, hhx_tf.shape)
+    self.assertAllClose(x, hhx_tf)
+
   @parameterized.parameters([2, 4, 8, 16])
   def test_output_same_as_simple_python_implementation(self, dim):
     """Tests result is identical to inefficient implementation using scipy."""