• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2016 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"""Weight broadcasting operations.
16
17In `tf.losses` and `tf.metrics`, we support limited weight broadcasting. This
18file includes operations for those broadcasting rules.
19"""
20
21from __future__ import absolute_import
22from __future__ import division
23from __future__ import print_function
24
25from tensorflow.python.framework import ops
26from tensorflow.python.framework import tensor_util
27from tensorflow.python.ops import array_ops
28from tensorflow.python.ops import control_flow_ops
29from tensorflow.python.ops import math_ops
30from tensorflow.python.ops import sets
31
32
33def _has_valid_dims(weights_shape, values_shape):
34  with ops.name_scope(
35      None, "has_invalid_dims", (weights_shape, values_shape)) as scope:
36    values_shape_2d = array_ops.expand_dims(values_shape, -1)
37    valid_dims = array_ops.concat(
38        (values_shape_2d, array_ops.ones_like(values_shape_2d)), axis=1)
39    weights_shape_2d = array_ops.expand_dims(weights_shape, -1)
40    invalid_dims = sets.set_difference(weights_shape_2d, valid_dims)
41    num_invalid_dims = array_ops.size(
42        invalid_dims.values, name="num_invalid_dims")
43    return math_ops.equal(0, num_invalid_dims, name=scope)
44
45
46def _has_valid_nonscalar_shape(
47    weights_rank, weights_shape, values_rank, values_shape):
48  with ops.name_scope(
49      None, "has_valid_nonscalar_shape",
50      (weights_rank, weights_shape, values_rank, values_shape)) as scope:
51    is_same_rank = math_ops.equal(
52        values_rank, weights_rank, name="is_same_rank")
53    return control_flow_ops.cond(
54        is_same_rank,
55        lambda: _has_valid_dims(weights_shape, values_shape),
56        lambda: is_same_rank,
57        name=scope)
58
59
60_ASSERT_BROADCASTABLE_ERROR_PREFIX = "weights can not be broadcast to values."
61
62
63def assert_broadcastable(weights, values):
64  """Asserts `weights` can be broadcast to `values`.
65
66  In `tf.losses` and `tf.metrics`, we support limited weight broadcasting. We
67  let weights be either scalar, or the same rank as the target values, with each
68  dimension either 1, or the same as the corresponding values dimension.
69
70  Args:
71    weights: `Tensor` of weights.
72    values: `Tensor` of values to which weights are applied.
73
74  Returns:
75    `Operation` raising `InvalidArgumentError` if `weights` has incorrect shape.
76    `no_op` if static checks determine `weights` has correct shape.
77
78  Raises:
79    ValueError:  If static checks determine `weights` has incorrect shape.
80  """
81  with ops.name_scope(None, "assert_broadcastable", (weights, values)) as scope:
82    with ops.name_scope(None, "weights", (weights,)) as weights_scope:
83      weights = ops.convert_to_tensor(weights, name=weights_scope)
84      weights_shape = array_ops.shape(weights, name="shape")
85      weights_rank = array_ops.rank(weights, name="rank")
86    weights_rank_static = tensor_util.constant_value(weights_rank)
87
88    with ops.name_scope(None, "values", (values,)) as values_scope:
89      values = ops.convert_to_tensor(values, name=values_scope)
90      values_shape = array_ops.shape(values, name="shape")
91      values_rank = array_ops.rank(values, name="rank")
92    values_rank_static = tensor_util.constant_value(values_rank)
93
94    # Try static checks.
95    if weights_rank_static is not None and values_rank_static is not None:
96      if weights_rank_static == 0:
97        return control_flow_ops.no_op(name="static_scalar_check_success")
98      if weights_rank_static != values_rank_static:
99        raise ValueError(
100            "%s values.rank=%s. weights.rank=%s."
101            " values.shape=%s. weights.shape=%s." % (
102                _ASSERT_BROADCASTABLE_ERROR_PREFIX, values_rank_static,
103                weights_rank_static, values.shape, weights.shape))
104      weights_shape_static = tensor_util.constant_value(weights_shape)
105      values_shape_static = tensor_util.constant_value(values_shape)
106      if weights_shape_static is not None and values_shape_static is not None:
107        # Sanity check, this should always be true since we checked rank above.
108        ndims = len(values_shape_static)
109        assert ndims == len(weights_shape_static)
110
111        for i in range(ndims):
112          if weights_shape_static[i] not in (1, values_shape_static[i]):
113            raise ValueError(
114                "%s Mismatch at dim %s. values.shape=%s weights.shape=%s." % (
115                    _ASSERT_BROADCASTABLE_ERROR_PREFIX, i, values_shape_static,
116                    weights_shape_static))
117        return control_flow_ops.no_op(name="static_dims_check_success")
118
119    # Dynamic checks.
120    is_scalar = math_ops.equal(0, weights_rank, name="is_scalar")
121    data = (
122        _ASSERT_BROADCASTABLE_ERROR_PREFIX,
123        "weights.shape=", weights.name, weights_shape,
124        "values.shape=", values.name, values_shape,
125        "is_scalar=", is_scalar,
126    )
127    is_valid_shape = control_flow_ops.cond(
128        is_scalar,
129        lambda: is_scalar,
130        lambda: _has_valid_nonscalar_shape(  # pylint: disable=g-long-lambda
131            weights_rank, weights_shape, values_rank, values_shape),
132        name="is_valid_shape")
133    return control_flow_ops.Assert(is_valid_shape, data, name=scope)
134
135
136def broadcast_weights(weights, values):
137  """Broadcast `weights` to the same shape as `values`.
138
139  This returns a version of `weights` following the same broadcast rules as
140  `mul(weights, values)`, but limited to the weights shapes allowed by
141  `assert_broadcastable`. When computing a weighted average, use this function
142  to broadcast `weights` before summing them; e.g.,
143  `reduce_sum(w * v) / reduce_sum(_broadcast_weights(w, v))`.
144
145  Args:
146    weights: `Tensor` whose shape is broadcastable to `values` according to the
147      rules of `assert_broadcastable`.
148    values: `Tensor` of any shape.
149
150  Returns:
151    `weights` broadcast to `values` shape according to the rules of
152      `assert_broadcastable`.
153  """
154  with ops.name_scope(None, "broadcast_weights", (weights, values)) as scope:
155    values = ops.convert_to_tensor(values, name="values")
156    weights = ops.convert_to_tensor(
157        weights, dtype=values.dtype.base_dtype, name="weights")
158
159    # Try static check for exact match.
160    weights_shape = weights.get_shape()
161    values_shape = values.get_shape()
162    if (weights_shape.is_fully_defined() and
163        values_shape.is_fully_defined() and
164        weights_shape.is_compatible_with(values_shape)):
165      return weights
166
167    with ops.control_dependencies((assert_broadcastable(weights, values),)):
168      return math_ops.multiply(
169          weights, array_ops.ones_like(values), name=scope)
170