1# Copyright 2017 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"""Utility ops shared across tf.contrib.signal."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import fractions
22
23from tensorflow.python.framework import ops
24from tensorflow.python.framework import tensor_util
25from tensorflow.python.ops import array_ops
26from tensorflow.python.ops import control_flow_ops
27from tensorflow.python.ops import math_ops
28
29
30def gcd(a, b, name=None):
31  """Returns the greatest common divisor via Euclid's algorithm.
32
33  Args:
34    a: The dividend. A scalar integer `Tensor`.
35    b: The divisor. A scalar integer `Tensor`.
36    name: An optional name for the operation.
37
38  Returns:
39    A scalar `Tensor` representing the greatest common divisor between `a` and
40    `b`.
41
42  Raises:
43    ValueError: If `a` or `b` are not scalar integers.
44  """
45  with ops.name_scope(name, 'gcd', [a, b]):
46    a = ops.convert_to_tensor(a)
47    b = ops.convert_to_tensor(b)
48
49    a.shape.assert_has_rank(0)
50    b.shape.assert_has_rank(0)
51
52    if not a.dtype.is_integer:
53      raise ValueError('a must be an integer type. Got: %s' % a.dtype)
54    if not b.dtype.is_integer:
55      raise ValueError('b must be an integer type. Got: %s' % b.dtype)
56
57    # TPU requires static shape inference. GCD is used for subframe size
58    # computation, so we should prefer static computation where possible.
59    const_a = tensor_util.constant_value(a)
60    const_b = tensor_util.constant_value(b)
61    if const_a is not None and const_b is not None:
62      return ops.convert_to_tensor(fractions.gcd(const_a, const_b))
63
64    cond = lambda _, b: math_ops.greater(b, array_ops.zeros_like(b))
65    body = lambda a, b: [b, math_ops.mod(a, b)]
66    a, b = control_flow_ops.while_loop(cond, body, [a, b], back_prop=False)
67    return a
68