def _tf_while_stmt(test, body, get_state, set_state, symbol_names, opts):
"""Overload of while_stmt that stages a TF while_stmt."""
init_vars = get_state()
orig_init_vars = init_vars
nulls = tuple(_is_none_or_undef(v) for v in init_vars)
if any(nulls):
shape_invars_by_init_vals = {
id(v): i for v, i in opts.get('shape_invariants', ())
}
shape_invariants = tuple(
shape_invars_by_init_vals.get(id(v), None) for v in orig_init_vars)
(require_one_iteration, init_vars,
extra_shape_invariants) = _try_handling_undefineds(body, get_state,
set_state, init_vars,
nulls, shape_invariants,
symbol_names)
else:
require_one_iteration = False
if require_one_iteration:
merged_shape_invariants = dict(shape_invars_by_init_vals)
# This has two roles:
# 1. Shape invariants are remapped from the old init vars to the new ones.
# 2. Any new shape invariants created by the init vars are kept, but only
# if the user didn't already specified some.
for v, nv, ni in zip(orig_init_vars, init_vars, extra_shape_invariants):
merged_invariant = merged_shape_invariants.get(id(v), ni)
if merged_invariant is not None:
merged_shape_invariants[id(nv)] = merged_invariant
merged_shape_invariants = tuple((nv, merged_shape_invariants[id(nv)])
for nv in init_vars
if id(nv) in merged_shape_invariants)
if merged_shape_invariants:
opts = dict(**opts)
opts['shape_invariants'] = merged_shape_invariants
def aug_test(*loop_vars):
if require_one_iteration:
loop_vars = loop_vars[1:]
set_state(loop_vars)
return _verify_tf_condition(test(), 'while loop')
def aug_body(*loop_vars):
if require_one_iteration:
loop_vars = loop_vars[1:]
set_state(loop_vars)
body()
new_loop_vars = get_state()
_verify_tf_loop_vars(
init_vars, loop_vars, new_loop_vars, symbol_names, opts)
if require_one_iteration:
new_loop_vars = (True,) + new_loop_vars
return new_loop_vars
if 'shape_invariants' in opts:
opts['shape_invariants'] = _shape_invariants_mapping_to_positional_list(
opts['shape_invariants'], init_vars)
while_loop_opts = dict(opts)
while_loop_opts.pop('iterate_names', None)
# Non-v2 while_loop unpacks the results when there is only one return value.
# This enforces consistency across versions.
while_loop_opts['return_same_structure'] = True
if require_one_iteration:
aug_init_vars = (False,) + init_vars
if 'shape_invariants' in while_loop_opts:
while_loop_opts['shape_invariants'] = (
(None,) + while_loop_opts['shape_invariants'])
else:
aug_init_vars = init_vars
final_loop_vars = control_flow_ops.while_loop(
aug_test, aug_body, aug_init_vars, **while_loop_opts)
if require_one_iteration:
with ops.control_dependencies([
control_flow_ops.Assert(final_loop_vars[0], [
_runtime_zero_iterations_errmsg(symbol_names, nulls, orig_init_vars)
])
]):
final_loop_vars = nest.map_structure(
lambda v: (array_ops.identity(v) if tensor_util.is_tf_type(v) else v),
final_loop_vars[1:],
)
set_state(final_loop_vars)
def step_fn(data):
assert_op = control_flow_ops.Assert(
math_ops.less_equal(math_ops.reduce_max(data), 100.), [data])
with ops.control_dependencies([assert_op]):
return math_ops.square(data)
def assert_greater(x):
assert_op = control_flow_ops.Assert(math_ops.greater(x, -1), [x])
with ops.control_dependencies([assert_op]):
return x
def squeeze(input, axis=None, name=None): # pylint: disable=redefined-builtin
"""Ragged compatible squeeze.
If `input` is a `tf.Tensor`, then this calls `tf.squeeze`.
If `input` is a `tf.RaggedTensor`, then this operation takes `O(N)` time,
where `N` is the number of elements in the squeezed dimensions.
Args:
input: A potentially ragged tensor. The input to squeeze.
axis: An optional list of ints. Defaults to `None`. If the `input` is
ragged, it only squeezes the dimensions listed. It fails if `input` is
ragged and axis is []. If `input` is not ragged it calls tf.squeeze. Note
that it is an error to squeeze a dimension that is not 1. It must be in
the range of [-rank(input), rank(input)).
name: A name for the operation (optional).
Returns:
A potentially ragged tensor. Contains the same data as input,
but has one or more dimensions of size 1 removed.
"""
with ops.name_scope(name, 'RaggedSqueeze', [input]):
input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input)
if isinstance(input, ops.Tensor):
return array_ops.squeeze(input, axis, name)
if axis is None:
raise ValueError('Ragged.squeeze must have an axis argument.')
if isinstance(axis, int):
axis = [axis]
elif ((not isinstance(axis, (list, tuple))) or
(not all(isinstance(d, int) for d in axis))):
raise TypeError('Axis must be a list or tuple of integers.')
dense_dims = []
ragged_dims = []
# Normalize all the dims in axis to be positive
axis = [ragged_util.get_positive_axis(d, input.shape.ndims) for d in axis]
for dim in axis:
if dim > input.ragged_rank:
dense_dims.append(dim - input.ragged_rank)
else:
ragged_dims.append(dim)
# Make sure the specified ragged dimensions are squeezable.
assertion_list = []
scalar_tensor_one = constant_op.constant(1, dtype=input.row_splits.dtype)
for i, r in enumerate(input.nested_row_lengths()):
if i + 1 in ragged_dims:
assertion_list.append(
control_flow_ops.Assert(
math_ops.reduce_all(math_ops.equal(r, scalar_tensor_one)),
['the given axis (axis = %d) is not squeezable!' % (i + 1)]))
if 0 in ragged_dims:
scalar_tensor_two = constant_op.constant(2, dtype=dtypes.int32)
assertion_list.append(
control_flow_ops.Assert(
math_ops.equal(
array_ops.size(input.row_splits), scalar_tensor_two),
['the given axis (axis = 0) is not squeezable!']))
# Till now, we are sure that the ragged dimensions are squeezable.
squeezed_rt = None
squeezed_rt = control_flow_ops.with_dependencies(assertion_list,
input.flat_values)
if dense_dims:
# Gives error if the dense dimension is not squeezable.
squeezed_rt = array_ops.squeeze(squeezed_rt, dense_dims)
remaining_row_splits = []
remaining_row_splits = list()
for i, row_split in enumerate(input.nested_row_splits):
# each row_splits tensor is for dimension #(i+1) .
if (i + 1) not in ragged_dims:
remaining_row_splits.append(row_split)
# Take care of the first row if it is to be squeezed.
if remaining_row_splits and 0 in ragged_dims:
remaining_row_splits.pop(0)
squeezed_rt = RaggedTensor.from_nested_row_splits(squeezed_rt,
remaining_row_splits)
# Corner case: when removing all the ragged dimensions and the output is
# a scalar tensor e.g. ragged.squeeze(ragged.constant([[[1]]])).
if set(range(0, input.ragged_rank + 1)).issubset(set(ragged_dims)):
squeezed_rt = array_ops.squeeze(squeezed_rt, [0], name)
return squeezed_rt
def guarded_body(i, s):
with ops.control_dependencies(
(control_flow_ops.Assert(i < 3, (i, )), )):
return s + i,
def fn(i):
error = control_flow_ops.Assert((i % 2) == 1, ["Error"])
with ops.control_dependencies([error]):
return v.read_value()
def get_weights_and_check_match_logits(
features, weight_column, logits, allow_per_logit_weights=False):
"""Fetches weights from features and checks that the shape matches logits.
Consider logits of shape [D0, D1, ... DN, logits_dimension]. Weights shape
can be either:
* [D0, D1, ... DN, logits_dimension] if `allow_per_logit_weights=True`.
* [D0, D1, ... DN, 1]
* [D0, D1, ... DN]: In this case, weights is reshaped into
[D0, D1, ... DN, 1] to work with weight broadcasting rules.
Args:
features: The features dict that contains weights.
weight_column: The weight column. If not given, this method returns 1.
logits: logits Tensor.
allow_per_logit_weights: Boolean. Whether we allow weights along the logits
dimension, namely shape `[D0, D1, ... DN, logits_dimension]`.
Returns:
Validated and reshaped weights Tensor.
Raises:
ValueError: If the weights `Tensor` cannot be cast into float.
"""
if allow_per_logit_weights:
err_msg = (
'weights shape must be [D0, D1, ... DN], [D0, D1, ... DN, 1] or '
'[D0, D1, ... DN, logits_dimension]')
else:
err_msg = (
'weights shape must be [D0, D1, ... DN] or [D0, D1, ... DN, 1]')
with ops.name_scope(
'weights', values=tuple(six.itervalues(features)) + (logits,)) as scope:
# Fetch the weights.
if weight_column is None:
return 1.
# TODO(b/117839674): update feature_column
if isinstance(weight_column, six.string_types):
weight_column = feature_column_lib.numeric_column(
key=weight_column, shape=(1,))
if not isinstance(weight_column,
(feature_column_lib.NumericColumn, _NumericColumn)):
raise TypeError('Weight column must be either a string or NumericColumn.'
' Given type: {}.'.format(type(weight_column)))
weights = weight_column._get_dense_tensor( # pylint: disable=protected-access
_LazyBuilder(features))
if not (weights.dtype.is_floating or weights.dtype.is_integer):
raise ValueError('Weight column should be castable to float. '
'Given dtype: {}'.format(weights.dtype))
weights = math_ops.cast(weights, name='weights', dtype=dtypes.float32)
# Validate the weights shape.
# Eager mode.
if context.executing_eagerly():
weights_shape = weights._shape_tuple() # pylint: disable=protected-access
logits_shape = logits._shape_tuple() # pylint: disable=protected-access
weights_rank = weights._rank() # pylint: disable=protected-access
logits_rank = logits._rank() # pylint: disable=protected-access
if (weights_rank is not None and logits_rank is not None
and weights_rank == logits_rank - 1):
if logits_shape[:-1] != weights_shape:
raise ValueError(
'{}, logits_shape: {}. weights_shape: {}.'
.format(err_msg, logits_shape, weights_shape))
return array_ops.expand_dims(weights, -1, name=scope)
supported_weights_shape = logits_shape[:-1] + (1,)
if allow_per_logit_weights:
if (logits_shape != weights_shape
and supported_weights_shape != weights_shape):
raise ValueError(
'{}, logits_shape: {}. weights_shape: {}.'
.format(err_msg, logits_shape, weights_shape))
else:
if supported_weights_shape != weights_shape:
raise ValueError(
'{}, logits_shape: {}. weights_shape: {}.'
.format(err_msg, logits_shape, weights_shape))
return weights
# Graph mode.
weights_shape = array_ops.shape(weights, name='weights_shape')
logits_shape = array_ops.shape(logits, name='logits_shape')
if (weights.shape.ndims is not None and logits.shape.ndims is not None and
weights.shape.ndims == logits.shape.ndims - 1):
assert_dimension = check_ops.assert_equal(
logits_shape[:-1], weights_shape, message=err_msg,
data=['logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
with ops.control_dependencies([assert_dimension]):
return array_ops.expand_dims(weights, -1, name=scope)
supported_weights_shape = array_ops.concat([logits_shape[:-1], [1]], axis=0)
if allow_per_logit_weights:
condition = math_ops.reduce_any(
[math_ops.reduce_all(math_ops.equal(logits_shape, weights_shape)),
math_ops.reduce_all(math_ops.equal(
supported_weights_shape, weights_shape))])
assert_dimension = control_flow_ops.Assert(
condition=condition,
data=[err_msg, 'logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
else:
assert_dimension = check_ops.assert_equal(
supported_weights_shape, weights_shape, message=err_msg,
data=['logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
with ops.control_dependencies([assert_dimension]):
return array_ops.identity(weights, name=scope)
def assert_near(x,
y,
rtol=None,
atol=None,
data=None,
summarize=None,
message=None,
name=None):
"""Assert the condition `x` and `y` are close element-wise.
Example of adding a dependency to an operation:
```python
with tf.control_dependencies([tf.assert_near(x, y)]):
output = tf.reduce_sum(x)
```
This condition holds if for every pair of (possibly broadcast) elements
`x[i]`, `y[i]`, we have
```tf.abs(x[i] - y[i]) <= atol + rtol * tf.abs(y[i])```.
If both `x` and `y` are empty, this is trivially satisfied.
The default `atol` and `rtol` is `10 * eps`, where `eps` is the smallest
representable positive number such that `1 + eps != eps`. This is about
`1.2e-6` in `32bit`, `2.22e-15` in `64bit`, and `0.00977` in `16bit`.
See `numpy.finfo`.
Args:
x: Float or complex `Tensor`.
y: Float or complex `Tensor`, same `dtype` as, and broadcastable to, `x`.
rtol: `Tensor`. Same `dtype` as, and broadcastable to, `x`.
The relative tolerance. Default is `10 * eps`.
atol: `Tensor`. Same `dtype` as, and broadcastable to, `x`.
The absolute tolerance. Default is `10 * eps`.
data: The tensors to print out if the condition is False. Defaults to
error message and first few entries of `x`, `y`.
summarize: Print this many entries of each tensor.
message: A string to prefix to the default message.
name: A name for this operation (optional). Defaults to "assert_near".
Returns:
Op that raises `InvalidArgumentError` if `x` and `y` are not close enough.
@compatibility(numpy)
Similar to `numpy.assert_allclose`, except tolerance depends on data type.
This is due to the fact that `TensorFlow` is often used with `32bit`, `64bit`,
and even `16bit` data.
@end_compatibility
"""
message = message or ''
with ops.name_scope(name, 'assert_near', [x, y, rtol, atol, data]):
x = ops.convert_to_tensor(x, name='x')
y = ops.convert_to_tensor(y, name='y', dtype=x.dtype)
eps = np.finfo(x.dtype.as_numpy_dtype).eps
rtol = 10 * eps if rtol is None else rtol
atol = 10 * eps if atol is None else atol
rtol = ops.convert_to_tensor(rtol, name='rtol', dtype=x.dtype)
atol = ops.convert_to_tensor(atol, name='atol', dtype=x.dtype)
if context.in_eager_mode():
x_name = _shape_and_dtype_str(x)
y_name = _shape_and_dtype_str(y)
else:
x_name = x.name
y_name = y.name
if data is None:
data = [
message,
'x and y not equal to tolerance rtol = %s, atol = %s' %
(rtol, atol),
'x (%s) = ' % x_name, x,
'y (%s) = ' % y_name, y
]
tol = atol + rtol * math_ops.abs(y)
diff = math_ops.abs(x - y)
condition = math_ops.reduce_all(math_ops.less(diff, tol))
return control_flow_ops.Assert(condition, data, summarize=summarize)
def stratified_sample(tensors,
labels,
target_probs,
batch_size,
init_probs=None,
enqueue_many=False,
queue_capacity=16,
threads_per_queue=1,
name=None):
"""Stochastically creates batches based on per-class probabilities.
This method discards examples. Internally, it creates one queue to amortize
the cost of disk reads, and one queue to hold the properly-proportioned
batch.
Args:
tensors: List of tensors for data. All tensors are either one item or a
batch, according to enqueue_many.
labels: Tensor for label of data. Label is a single integer or a batch,
depending on enqueue_many. It is not a one-hot vector.
target_probs: Target class proportions in batch. An object whose type has a
registered Tensor conversion function.
batch_size: Size of batch to be returned.
init_probs: Class proportions in the data. An object whose type has a
registered Tensor conversion function, or `None` for estimating the
initial distribution.
enqueue_many: Bool. If true, interpret input tensors as having a batch
dimension.
queue_capacity: Capacity of the large queue that holds input examples.
threads_per_queue: Number of threads for the large queue that holds input
examples and for the final queue with the proper class proportions.
name: Optional prefix for ops created by this function.
Raises:
ValueError: enqueue_many is True and labels doesn't have a batch
dimension, or if enqueue_many is False and labels isn't a scalar.
ValueError: enqueue_many is True, and batch dimension on data and labels
don't match.
ValueError: if probs don't sum to one.
ValueError: if a zero initial probability class has a nonzero target
probability.
TFAssertion: if labels aren't integers in [0, num classes).
Returns:
(data_batch, label_batch), where data_batch is a list of tensors of the same
length as `tensors`
Example:
# Get tensor for a single data and label example.
data, label = data_provider.Get(['data', 'label'])
# Get stratified batch according to per-class probabilities.
target_probs = [...distribution you want...]
[data_batch], labels = tf.contrib.training.stratified_sample(
[data], label, target_probs)
# Run batch through network.
...
"""
with ops.name_scope(name, 'stratified_sample', tensors + [labels]):
tensor_list = ops.convert_n_to_tensor_or_indexed_slices(tensors)
labels = ops.convert_to_tensor(labels)
target_probs = ops.convert_to_tensor(target_probs,
dtype=dtypes.float32)
# Reduce the case of a single example to that of a batch of size 1.
if not enqueue_many:
tensor_list = [
array_ops.expand_dims(tensor, 0) for tensor in tensor_list
]
labels = array_ops.expand_dims(labels, 0)
# If `init_probs` is `None`, set up online estimation of data distribution.
if init_probs is None:
# We use `target_probs` to get the number of classes, so its shape must be
# fully defined at graph construction time.
target_probs.get_shape().assert_is_fully_defined()
init_probs = _estimate_data_distribution(
labels,
target_probs.get_shape().num_elements())
else:
init_probs = ops.convert_to_tensor(init_probs,
dtype=dtypes.float32)
# Validate that input is consistent.
tensor_list, labels, [init_probs, target_probs
] = _verify_input(tensor_list, labels,
[init_probs, target_probs])
# Check that all zero initial probabilities also have zero target
# probabilities.
assert_op = control_flow_ops.Assert(
math_ops.reduce_all(
math_ops.logical_or(math_ops.not_equal(init_probs, 0),
math_ops.equal(target_probs, 0))),
[
'All classes with zero initial probability must also have zero target '
'probability: ', init_probs, target_probs
])
init_probs = control_flow_ops.with_dependencies([assert_op],
init_probs)
# Calculate acceptance sampling probabilities.
accept_probs = _calculate_acceptance_probabilities(
init_probs, target_probs)
proportion_rejected = math_ops.reduce_sum(
(1 - accept_probs) * init_probs)
accept_probs = control_flow_ops.cond(
math_ops.less(proportion_rejected, .5),
lambda: accept_probs,
lambda: logging_ops.Print( # pylint: disable=g-long-lambda
accept_probs, [accept_probs],
message='Proportion of examples rejected by sampler is high.',
first_n=10))
# Make a single queue to hold input examples. Reshape output so examples
# don't have singleton batch dimension.
batched = input_ops.batch(tensor_list + [labels],
batch_size=1,
num_threads=threads_per_queue,
capacity=queue_capacity,
enqueue_many=True)
val_list = [array_ops.squeeze(x, [0]) for x in batched[:-1]]
label = array_ops.squeeze(batched[-1], [0])
# Set up second queue containing batches that have the desired class
# proportions.
cur_prob = array_ops.gather(accept_probs, label)
batched = input_ops.maybe_batch(
val_list + [label],
keep_input=random_ops.random_uniform([]) < cur_prob,
batch_size=batch_size,
num_threads=threads_per_queue)
return batched[:-1], batched[-1]
def _tf_while_stmt(test, body, get_state, set_state, symbol_names, opts):
"""Overload of while_stmt that stages a TF while_stmt."""
init_vars = get_state()
orig_init_vars = init_vars
nulls = tuple(_is_none_or_undef(v) for v in init_vars)
if any(nulls):
require_one_iteration, init_vars = _try_handling_undefineds(
body, get_state, set_state, init_vars, nulls, symbol_names)
else:
require_one_iteration = False
def aug_test(*loop_vars):
if require_one_iteration:
loop_vars = loop_vars[1:]
set_state(loop_vars)
return test()
def aug_body(*loop_vars):
if require_one_iteration:
loop_vars = loop_vars[1:]
set_state(loop_vars)
body()
new_loop_vars = get_state()
_verify_tf_loop_vars(
init_vars, loop_vars, new_loop_vars, symbol_names, opts)
if require_one_iteration:
new_loop_vars = (True,) + new_loop_vars
return new_loop_vars
if 'shape_invariants' in opts:
opts['shape_invariants'] = _shape_invariants_mapping_to_positional_list(
opts['shape_invariants'], init_vars)
while_loop_opts = dict(opts)
while_loop_opts.pop('iterate_names', None)
# Non-v2 while_loop unpacks the results when there is only one return value.
# This enforces consistency across versions.
while_loop_opts['return_same_structure'] = True
if require_one_iteration:
aug_init_vars = (False,) + init_vars
else:
aug_init_vars = init_vars
final_loop_vars = control_flow_ops.while_loop(
aug_test, aug_body, aug_init_vars, **while_loop_opts)
if require_one_iteration:
with ops.control_dependencies([
control_flow_ops.Assert(final_loop_vars[0], [
_runtime_zero_iterations_errmsg(symbol_names, nulls, orig_init_vars)
])
]):
final_loop_vars = tuple(
array_ops.identity(v) for v in final_loop_vars[1:])
set_state(final_loop_vars)
def assert_equal(x, y, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x == y` holds element-wise.
Example of adding a dependency to an operation:
```python
with tf.control_dependencies([tf.assert_equal(x, y)]):
output = tf.reduce_sum(x)
```
This condition holds if for every pair of (possibly broadcast) elements
`x[i]`, `y[i]`, we have `x[i] == y[i]`.
If both `x` and `y` are empty, this is trivially satisfied.
Args:
x: Numeric `Tensor`.
y: Numeric `Tensor`, same dtype as and broadcastable to `x`.
data: The tensors to print out if the condition is False. Defaults to
error message and first few entries of `x`, `y`.
summarize: Print this many entries of each tensor.
message: A string to prefix to the default message.
name: A name for this operation (optional). Defaults to "assert_equal".
Returns:
Op that raises `InvalidArgumentError` if `x == y` is False.
@compatibility{eager} returns None
Raises:
InvalidArgumentError if the check can be performed immediately and
`x == y` is False. The check can be performed immediately during
eager execution or if `x` and `y` are statically known.
"""
message = message or ''
with ops.name_scope(name, 'assert_equal', [x, y, data]):
x = ops.convert_to_tensor(x, name='x')
y = ops.convert_to_tensor(y, name='y')
if context.in_eager_mode():
eq = math_ops.equal(x, y)
condition = math_ops.reduce_all(eq)
if not condition:
# Prepare a message with first elements of x and y.
summary_msg = ''
# Default to printing 3 elements like control_flow_ops.Assert (used
# by graph mode) does.
summarize = 3 if summarize is None else summarize
if summarize:
# reshape((-1,)) is the fastest way to get a flat array view.
x_np = x.numpy().reshape((-1, ))
y_np = y.numpy().reshape((-1, ))
x_sum = min(x_np.size, summarize)
y_sum = min(y_np.size, summarize)
summary_msg = ('First %d elements of x:\n%s\n'
'First %d elements of y:\n%s\n' %
(x_sum, x_np[:x_sum], y_sum, y_np[:y_sum]))
# Get the values that actually differed and their indices.
mask = math_ops.logical_not(eq)
indices = array_ops.where(mask)
indices_np = indices.numpy()
x_vals = array_ops.boolean_mask(x, mask)
y_vals = array_ops.boolean_mask(y, mask)
summarize = min(summarize, indices_np.shape[0])
raise errors.InvalidArgumentError(
node_def=None,
op=None,
message=('%s\nCondition x == y did not hold.\n'
'Indices of first %s different values:\n%s\n'
'Corresponding x values:\n%s\n'
'Corresponding y values:\n%s\n'
'%s' %
(message or '', summarize, indices_np[:summarize],
x_vals.numpy().reshape(
(-1, ))[:summarize], y_vals.numpy().reshape(
(-1, ))[:summarize], summary_msg)))
return
if data is None:
data = [
message, 'Condition x == y did not hold element-wise:',
'x (%s) = ' % x.name, x,
'y (%s) = ' % y.name, y
]
condition = math_ops.reduce_all(math_ops.equal(x, y))
x_static = tensor_util.constant_value(x)
y_static = tensor_util.constant_value(y)
if x_static is not None and y_static is not None:
condition_static = (x_static == y_static).all()
_assert_static(condition_static, data)
return control_flow_ops.Assert(condition, data, summarize=summarize)
def _sampled_scattered_embedding_lookup(
params, values, dimension=None, sampled_candidates=None, hash_key=None,
name=None):
"""Looks up embeddings using parameter hashing for each value in `values`.
This method looks up selected embedding dimensions if `sampled_candidates` is
given, otherwise looks up all dimensions.
The i-th embedding component of a value v in `values` is found by retrieving
the weight whose index is a fingerprint of the pair (v,i).
The concept is explored as "feature hashing" for model compression in this
paper: http://arxiv.org/pdf/1504.04788.pdf
Feature hashing has the pleasant effect of allowing us to compute an embedding
without needing a pre-determined vocabulary, relieving some amount of process
complexity. It also allows for us to maintain embeddings for possibly
trillions of features with a fixed amount of memory.
Note that this is superior to out-of-vocabulary shared "hash buckets" in that
the embedding is extremely likely to be unique for each token as opposed to
being shared across probably-colliding tokens. The price is that we must
compute a hash once for each scalar in the token's embedding as opposed to
once per token.
If `params` is a list, it represents a partition of the embedding parameters.
Each tensor in the list should have the same length, except for the first ones
which may have an additional element. For instance 10 parameters can be
partitioned in 4 tensors with length `[3, 3, 2, 2]`.
Args:
params: A `Tensor`, `list` of `Tensors`, or `PartitionedVariable`.
Each tensor must be of rank 1 with fully-defined shape.
values: `Tensor` of values to be embedded with shape `[d0, ..., dn]`.
dimension: Embedding dimension. The user must specify either `dimension` or
`sampled_candidates`.
sampled_candidates: An optional `Tensor` of slice indices to keep along the
final dimension with shape `[d0, ..., dn, N]`. If given, `dimension` is
ignored. If `None`, looks up all candidates.
hash_key: Specify the hash_key that will be used by the `FingerprintCat64`
function to combine the crosses fingerprints on SparseFeatureCrossOp
(optional).
name: An optional name for this op.
Returns:
A `Tensor` with shape `[d0, ..., dn, dimension]`.
If `sampled_candidates` is given, the output shape is `[d0, ..., dn, N]`
Raises:
ValueError: if dimension is not positive or the partition size is invalid.
"""
if isinstance(params, variables.PartitionedVariable):
params = list(params)
if not isinstance(params, list):
params = [params]
with ops.name_scope(name, "scattered_embedding_lookup",
params + [dimension, values]):
# Flatten the values
values_shape = array_ops.shape(values)
values = array_ops.reshape(values, [-1, 1])
if sampled_candidates is None:
if dimension is None:
raise ValueError(
"You must specify either dimension or sampled_candidates.")
if dimension <= 0:
raise ValueError("Dimension must be >0. Given is %d" % dimension)
sampled_candidates = array_ops.tile(array_ops.expand_dims(
math_ops.range(0, dimension), 0), array_ops.shape(values))
else:
dimension = array_ops.shape(sampled_candidates)[
math_ops.subtract(array_ops.rank(sampled_candidates), 1)]
sampled_candidates_shape = array_ops.shape(sampled_candidates)
dimension_tensor = array_ops.reshape(dimension, shape=[1,])
expected_shape = array_ops.concat([values_shape, dimension_tensor], 0)
with ops.control_dependencies([control_flow_ops.Assert(
math_ops.reduce_all(math_ops.equal(sampled_candidates_shape,
expected_shape)),
["The shape of sampled_candidates: ", sampled_candidates_shape,
" does not match the shape of values: ", values_shape])]):
# Flatten sampled_candidates, same way as values are flattened.
sampled_candidates = array_ops.reshape(sampled_candidates,
[-1, dimension])
num_partitions = len(params)
partition_sizes = []
for p in range(num_partitions):
shape = params[p].get_shape()
shape.assert_has_rank(1)
shape.assert_is_fully_defined()
partition_sizes.append(shape[0].value)
num_params = sum(partition_sizes) # Total number of parameters.
# Assert the size of each partition.
for p in range(num_partitions):
expected_size = (num_params - p - 1) // num_partitions + 1
if partition_sizes[p] != expected_size:
raise ValueError("Tensor %d in params has size %d, expected %d." %
(p, partition_sizes[p], expected_size))
# With two values v1 and v2 and 3 dimensions, we will cross
# [[0, 1, 2], [0, 1, 2]] with [[v1], [v2]].
tensors_to_cross = [sampled_candidates, values]
ids = sparse_feature_cross_op.sparse_feature_cross(
tensors_to_cross, hashed_output=True, num_buckets=num_params,
hash_key=hash_key)
ids = sparse_ops.sparse_tensor_to_dense(ids)
# No need to validate the indices since we have checked the params
# dimensions and we know the largest id.
result = embedding_ops.embedding_lookup(
params, ids, partition_strategy="div")
return array_ops.reshape(result,
array_ops.concat([values_shape, [dimension]], 0))
def _dopri5(func,
y0,
t,
rtol,
atol,
full_output=False,
first_step=None,
safety=0.9,
ifactor=10.0,
dfactor=0.2,
max_num_steps=1000,
name=None):
"""Solve an ODE for `odeint` using method='dopri5'."""
if first_step is None:
# at some point, we might want to switch to picking the step size
# automatically
first_step = 1.0
with ops.name_scope(
name, 'dopri5',
[y0, t, rtol, atol, safety, ifactor, dfactor, max_num_steps]) as scope:
first_step = ops.convert_to_tensor(first_step, dtype=t.dtype,
name='first_step')
safety = ops.convert_to_tensor(safety, dtype=t.dtype, name='safety')
ifactor = ops.convert_to_tensor(ifactor, dtype=t.dtype, name='ifactor')
dfactor = ops.convert_to_tensor(dfactor, dtype=t.dtype, name='dfactor')
max_num_steps = ops.convert_to_tensor(max_num_steps, dtype=dtypes.int32,
name='max_num_steps')
def adaptive_runge_kutta_step(rk_state, history, n_steps):
"""Take an adaptive Runge-Kutta step to integrate the ODE."""
y0, f0, _, t0, dt, interp_coeff = rk_state
with ops.name_scope('assertions'):
check_underflow = control_flow_ops.Assert(
t0 + dt > t0, ['underflow in dt', dt])
check_max_num_steps = control_flow_ops.Assert(
n_steps < max_num_steps, ['max_num_steps exceeded'])
check_numerics = control_flow_ops.Assert(
math_ops.reduce_all(math_ops.is_finite(abs(y0))),
['non-finite values in state `y`', y0])
with ops.control_dependencies(
[check_underflow, check_max_num_steps, check_numerics]):
y1, f1, y1_error, k = _runge_kutta_step(func, y0, f0, t0, dt)
with ops.name_scope('error_ratio'):
# We use the same approach as the dopri5 fortran code.
error_tol = atol + rtol * math_ops.maximum(abs(y0), abs(y1))
tensor_error_ratio = _abs_square(y1_error) / _abs_square(error_tol)
# Could also use reduce_maximum here.
error_ratio = math_ops.sqrt(math_ops.reduce_mean(tensor_error_ratio))
accept_step = error_ratio <= 1
with ops.name_scope('update/rk_state'):
# If we don't accept the step, the _RungeKuttaState will be useless
# (covering a time-interval of size 0), but that's OK, because in such
# cases we always immediately take another Runge-Kutta step.
y_next = control_flow_ops.cond(accept_step, lambda: y1, lambda: y0)
f_next = control_flow_ops.cond(accept_step, lambda: f1, lambda: f0)
t_next = control_flow_ops.cond(accept_step, lambda: t0 + dt, lambda: t0)
interp_coeff = control_flow_ops.cond(
accept_step,
lambda: _interp_fit_rk(y0, y1, k, dt),
lambda: interp_coeff)
dt_next = _optimal_step_size(dt, error_ratio, safety, ifactor, dfactor)
rk_state = _RungeKuttaState(
y_next, f_next, t0, t_next, dt_next, interp_coeff)
with ops.name_scope('update/history'):
history = _History(_ta_append(history.integrate_points, t0 + dt),
_ta_append(history.error_ratio, error_ratio))
return rk_state, history, n_steps + 1
def interpolate(solution, history, rk_state, i):
"""Interpolate through the next time point, integrating as necessary."""
with ops.name_scope('interpolate'):
rk_state, history, _ = control_flow_ops.while_loop(
lambda rk_state, *_: t[i] > rk_state.t1,
adaptive_runge_kutta_step,
(rk_state, history, 0),
name='integrate_loop')
y = _interp_evaluate(
rk_state.interp_coeff, rk_state.t0, rk_state.t1, t[i])
solution = solution.write(i, y)
return solution, history, rk_state, i + 1
assert_increasing = control_flow_ops.Assert(
math_ops.reduce_all(t[1:] > t[:-1]),
['`t` must be monotonic increasing'])
with ops.control_dependencies([assert_increasing]):
num_times = array_ops.size(t)
solution = tensor_array_ops.TensorArray(
y0.dtype, size=num_times).write(0, y0)
history = _History(
integrate_points=tensor_array_ops.TensorArray(
t.dtype, size=0, dynamic_size=True),
error_ratio=tensor_array_ops.TensorArray(
rtol.dtype, size=0, dynamic_size=True))
rk_state = _RungeKuttaState(
y0, func(y0, t[0]), t[0], t[0], first_step, interp_coeff=[y0] * 5)
solution, history, _, _ = control_flow_ops.while_loop(
lambda _, __, ___, i: i < num_times,
interpolate,
(solution, history, rk_state, 1),
name='interpolate_loop')
y = solution.stack(name=scope)
y.set_shape(t.get_shape().concatenate(y0.get_shape()))
if not full_output:
return y
else:
integrate_points = history.integrate_points.stack()
info_dict = {'num_func_evals': 6 * array_ops.size(integrate_points) + 1,
'integrate_points': integrate_points,
'error_ratio': history.error_ratio.stack()}
return (y, info_dict)
def _training_examples_and_variables():
"""Returns dictionaries for training examples and variables."""
batch_size = targets.get_shape()[0]
# Iterate over all feature columns and create appropriate lists for dense
# and sparse features as well as dense and sparse weights (variables) for
# SDCA.
# TODO(sibyl-vie3Poto): Reshape variables stored as values in column_to_variables
# dict as 1-dimensional tensors.
dense_features, sparse_features, sparse_feature_with_values = [], [], []
dense_feature_weights = []
sparse_feature_weights, sparse_feature_with_values_weights = [], []
for column in sorted(columns_to_variables.keys(),
key=lambda x: x.key):
transformed_tensor = features[column]
if isinstance(column, layers.feature_column._RealValuedColumn): # pylint: disable=protected-access
# A real-valued column corresponds to a dense feature in SDCA. A
# transformed tensor corresponding to a RealValuedColumn should have
# rank at most 2. In order to be passed to SDCA, its rank needs to be
# exactly 2 (i.e., its shape should be [batch_size, column.dim]).
check_rank_op = control_flow_ops.Assert(
math_ops.less_equal(array_ops.rank(transformed_tensor),
2),
['transformed_tensor should have rank at most 2.'])
# Reshape to [batch_size, dense_column_dimension].
with ops.control_dependencies([check_rank_op]):
transformed_tensor = array_ops.reshape(
transformed_tensor,
[array_ops.shape(transformed_tensor)[0], -1])
dense_features.append(transformed_tensor)
# For real valued columns, the variables list contains exactly one
# element.
dense_feature_weights.append(
columns_to_variables[column][0])
elif isinstance(column,
layers.feature_column._BucketizedColumn): # pylint: disable=protected-access
# A bucketized column corresponds to a sparse feature in SDCA. The
# bucketized feature is "sparsified" for SDCA by converting it to a
# SparseFeatureColumn representing the one-hot encoding of the
# bucketized feature.
#
# TODO(sibyl-vie3Poto): Explore whether it is more efficient to translate a
# bucketized feature column to a dense feature in SDCA. This will
# likely depend on the number of buckets.
dense_bucket_tensor = column._to_dnn_input_layer(
transformed_tensor) # pylint: disable=protected-access
sparse_feature_column = _dense_tensor_to_sparse_feature_column(
dense_bucket_tensor)
sparse_feature_with_values.append(sparse_feature_column)
# If a partitioner was used during variable creation, we will have a
# list of Variables here larger than 1.
vars_to_append = columns_to_variables[column][0]
if len(columns_to_variables[column]) > 1:
vars_to_append = columns_to_variables[column]
sparse_feature_with_values_weights.append(vars_to_append)
elif isinstance(
column,
(
layers.feature_column._WeightedSparseColumn, # pylint: disable=protected-access
layers.feature_column._CrossedColumn, # pylint: disable=protected-access
layers.feature_column._SparseColumn)): # pylint: disable=protected-access
if isinstance(column,
layers.feature_column._WeightedSparseColumn): # pylint: disable=protected-access
id_tensor = column.id_tensor(transformed_tensor)
weight_tensor = array_ops.reshape(
column.weight_tensor(transformed_tensor).values,
[-1])
else:
id_tensor = transformed_tensor
weight_tensor = array_ops.ones(
[array_ops.shape(id_tensor.indices)[0]],
dtypes.float32)
example_ids = array_ops.reshape(id_tensor.indices[:, 0],
[-1])
flat_ids = array_ops.reshape(id_tensor.values, [-1])
# Prune invalid IDs (< 0) from the flat_ids, example_ids, and
# weight_tensor. These can come from looking up an OOV entry in the
# vocabulary (default value being -1).
is_id_valid = math_ops.greater_equal(flat_ids, 0)
flat_ids = array_ops.boolean_mask(flat_ids, is_id_valid)
example_ids = array_ops.boolean_mask(
example_ids, is_id_valid)
weight_tensor = array_ops.boolean_mask(
weight_tensor, is_id_valid)
projection_length = math_ops.reduce_max(flat_ids) + 1
# project ids based on example ids so that we can dedup ids that
# occur multiple times for a single example.
projected_ids = projection_length * example_ids + flat_ids
# Remove any redundant ids.
ids, idx = array_ops.unique(projected_ids)
# Keep only one example id per duplicated ids.
example_ids_filtered = math_ops.unsorted_segment_min(
example_ids, idx,
array_ops.shape(ids)[0])
# reproject ids back feature id space.
reproject_ids = (ids -
projection_length * example_ids_filtered)
weights = array_ops.reshape(
math_ops.unsorted_segment_sum(weight_tensor, idx,
array_ops.shape(ids)[0]),
[-1])
sparse_feature_with_values.append(
SparseFeatureColumn(example_ids_filtered,
reproject_ids, weights))
# If a partitioner was used during variable creation, we will have a
# list of Variables here larger than 1.
vars_to_append = columns_to_variables[column][0]
if len(columns_to_variables[column]) > 1:
vars_to_append = columns_to_variables[column]
sparse_feature_with_values_weights.append(vars_to_append)
else:
raise ValueError(
'SDCAOptimizer does not support column type %s.' %
type(column).__name__)
example_weights = array_ops.reshape(
features[weight_column_name], shape=[
-1
]) if weight_column_name else array_ops.ones([batch_size])
example_ids = features[self._example_id_column]
sparse_feature_with_values.extend(sparse_features)
sparse_feature_with_values_weights.extend(sparse_feature_weights)
examples = dict(sparse_features=sparse_feature_with_values,
dense_features=dense_features,
example_labels=math_ops.to_float(
array_ops.reshape(targets, shape=[-1])),
example_weights=example_weights,
example_ids=example_ids)
sdca_variables = dict(
sparse_features_weights=sparse_feature_with_values_weights,
dense_features_weights=dense_feature_weights)
return examples, sdca_variables
def detokenize(self, token_ids):
r"""Convert a `Tensor` or `RaggedTensor` of wordpiece IDs to string-words.
>>> import pathlib
>>> pathlib.Path('/tmp/detok_vocab.txt').write_text(
... 'a b c ##a ##b ##c'.replace(' ', '\n'))
>>> wordpiece = WordpieceTokenizer('/tmp/detok_vocab.txt')
>>> token_ids = [[0, 4, 5, 2, 5, 5, 5]]
>>> wordpiece.detokenize(token_ids)
<tf.RaggedTensor [[b'abc', b'cccc']]>
The word pieces are joined along the innermost axis to make words. So the
result has the same rank as the input, but the innermost axis of the result
indexes words instead of word pieces.
The shape transformation is: `[..., wordpieces] => [..., words]`
When the input shape is `[..., words, wordpieces]` (like the output of
`WordpieceTokenizer.tokenize`) the result's shape is `[..., words, 1]`.
The additional ragged axis can be removed using `words.merge_dims(-2, -1)`.
Note: This method assumes wordpiece IDs are dense on the interval
`[0, vocab_size)`.
Args:
token_ids: A `RaggedTensor` or `Tensor` with an int dtype. Must have
`ndims >= 2`
Returns:
A `RaggedTensor` with dtype `string` and the rank as the input
`token_ids`.
"""
# If there are performance issues with this method or problems with lookup
# tables using sparse IDs see the notes in b/177610044.
vocab, ids = self._get_vocab_and_ids()
token_ids = ragged_tensor.convert_to_tensor_or_ragged_tensor(token_ids)
first_is_zero = math_ops.equal(ids[0], 0)
steps = ids[1:] - ids[:-1]
all_one_step = math_ops.reduce_all(math_ops.equal(steps, 1))
check = control_flow_ops.Assert(
first_is_zero & all_one_step,
data=[('`detokenize` only works with vocabulary tables where the '
'indices are dense on the interval `[0, vocab_size)`')])
with ops.control_dependencies([check]):
token_ids = math_ops.minimum(
token_ids,
# Limit the OOV buckets to a single index.
math_ops.cast(array_ops.size(vocab), token_ids.dtype))
# Add the unknown token at that index.
vocab = array_ops.concat([vocab, [self._unknown_token]], axis=0)
# Lookup the text tokens and join them along the innermost axis.
txt_tokens = array_ops.gather(vocab, token_ids)
# Ensure the input is Ragged.
if not isinstance(txt_tokens, RaggedTensor):
txt_tokens = RaggedTensor.from_tensor(txt_tokens)
# Join the tokens along the last axis.
words = string_ops.reduce_join_v2(txt_tokens, axis=-1, separator=' ')
# Collapse " ##" in all strings to make words.
words = string_ops.regex_replace(
words, ' ' + re.escape(self._suffix_indicator), '')
# Strip leading and trailing spaces.
words = string_ops.regex_replace(words, '^ +| +$', '')
# Split on spaces so the last axis is "words".
words = ragged_string_ops.string_split_v2(words, sep=' ')
return words
def _parse_single_example_raw(serialized,
names=None,
sparse_keys=None,
sparse_types=None,
dense_keys=None,
dense_types=None,
dense_defaults=None,
dense_shapes=None,
name=None):
"""Parses a single `Example` proto.
Args:
serialized: A scalar string Tensor, a single serialized Example.
See `_parse_example_raw` documentation for more details.
names: (Optional) A scalar string Tensor, the associated name.
See `_parse_example_raw` documentation for more details.
sparse_keys: See `_parse_example_raw` documentation for more details.
sparse_types: See `_parse_example_raw` documentation for more details.
dense_keys: See `_parse_example_raw` documentation for more details.
dense_types: See `_parse_example_raw` documentation for more details.
dense_defaults: See `_parse_example_raw` documentation for more details.
dense_shapes: See `_parse_example_raw` documentation for more details.
name: A name for this operation (optional).
Returns:
A `dict` mapping feature keys to `Tensor` and `SparseTensor` values.
Raises:
ValueError: if any feature is invalid.
"""
with ops.name_scope(name, "ParseSingleExample", [serialized, names]):
serialized = ops.convert_to_tensor(serialized)
serialized_shape = serialized.get_shape()
if serialized_shape.ndims is not None:
if serialized_shape.ndims != 0:
raise ValueError("Input serialized must be a scalar")
else:
serialized = control_flow_ops.with_dependencies(
[
control_flow_ops.Assert(math_ops.equal(
array_ops.rank(serialized),
0), ["Input serialized must be a scalar"],
name="SerializedIsScalar")
],
serialized,
name="SerializedDependencies")
serialized = array_ops.expand_dims(serialized, 0)
if names is not None:
names = ops.convert_to_tensor(names)
names_shape = names.get_shape()
if names_shape.ndims is not None:
if names_shape.ndims != 0:
raise ValueError("Input names must be a scalar")
else:
names = control_flow_ops.with_dependencies(
[
control_flow_ops.Assert(math_ops.equal(
array_ops.rank(names),
0), ["Input names must be a scalar"],
name="NamesIsScalar")
],
names,
name="NamesDependencies")
names = array_ops.expand_dims(names, 0)
outputs = _parse_example_raw(serialized,
names=names,
sparse_keys=sparse_keys,
sparse_types=sparse_types,
dense_keys=dense_keys,
dense_types=dense_types,
dense_defaults=dense_defaults,
dense_shapes=dense_shapes,
name=name)
if dense_keys is not None:
for d in dense_keys:
d_name = re.sub("[^A-Za-z0-9_.\\-/]", "_", d)
outputs[d] = array_ops.squeeze(outputs[d], [0],
name="Squeeze_%s" % d_name)
if sparse_keys is not None:
for s in sparse_keys:
s_name = re.sub("[^A-Za-z0-9_.\\-/]", "_", s)
outputs[s] = sparse_tensor.SparseTensor(
array_ops.slice(outputs[s].indices, [0, 1], [-1, -1],
name="Slice_Indices_%s" % s_name),
outputs[s].values,
array_ops.slice(outputs[s].dense_shape, [1], [-1],
name="Squeeze_Shape_%s" % s_name))
return outputs
def f(x):
control_flow_ops.Assert(x == 1, ['Wrong value'])
def assert_broadcastable(weights, values):
"""Asserts `weights` can be broadcast to `values`.
In `tf.losses` and `tf.metrics`, we support limited weight broadcasting. We
let weights be either scalar, or the same rank as the target values, with each
dimension either 1, or the same as the corresponding values dimension.
Args:
weights: `Tensor` of weights.
values: `Tensor` of values to which weights are applied.
Returns:
`Operation` raising `InvalidArgumentError` if `weights` has incorrect shape.
`no_op` if static checks determine `weights` has correct shape.
Raises:
ValueError: If static checks determine `weights` has incorrect shape.
"""
with ops.name_scope(None, "assert_broadcastable",
(weights, values)) as scope:
with ops.name_scope(None, "weights", (weights, )) as weights_scope:
weights = ops.convert_to_tensor(weights, name=weights_scope)
weights_shape = array_ops.shape(weights, name="shape")
weights_rank = array_ops.rank(weights, name="rank")
weights_rank_static = tensor_util.constant_value(weights_rank)
with ops.name_scope(None, "values", (values, )) as values_scope:
values = ops.convert_to_tensor(values, name=values_scope)
values_shape = array_ops.shape(values, name="shape")
values_rank = array_ops.rank(values, name="rank")
values_rank_static = tensor_util.constant_value(values_rank)
# Try static checks.
if weights_rank_static is not None and values_rank_static is not None:
if weights_rank_static == 0:
return control_flow_ops.no_op(
name="static_scalar_check_success")
if weights_rank_static != values_rank_static:
raise ValueError("%s values.rank=%s. weights.rank=%s." %
(_ASSERT_BROADCASTABLE_ERROR_PREFIX,
values_rank_static, weights_rank_static))
weights_shape_static = tensor_util.constant_value(weights_shape)
values_shape_static = tensor_util.constant_value(values_shape)
if weights_shape_static is not None and values_shape_static is not None:
# Sanity check, this should always be true since we checked rank above.
ndims = len(values_shape_static)
assert ndims == len(weights_shape_static)
for i in range(ndims):
if weights_shape_static[i] not in (1,
values_shape_static[i]):
raise ValueError(
"%s Mismatch at dim %s. values.shape=%s weights.shape=%s."
% (_ASSERT_BROADCASTABLE_ERROR_PREFIX, i,
values_shape_static, weights_shape_static))
return control_flow_ops.no_op(name="static_dims_check_success")
# Dynamic checks.
is_scalar = math_ops.equal(0, weights_rank, name="is_scalar")
data = (
_ASSERT_BROADCASTABLE_ERROR_PREFIX,
"weights.shape=",
weights.name,
weights_shape,
"values.shape=",
values.name,
values_shape,
"is_scalar=",
is_scalar,
)
is_valid_shape = control_flow_ops.cond(
is_scalar,
lambda: is_scalar,
lambda: _has_valid_nonscalar_shape( # pylint: disable=g-long-lambda
weights_rank, weights_shape, values_rank, values_shape),
name="is_valid_shape")
return control_flow_ops.Assert(is_valid_shape, data, name=scope)
def loop_fn(i): return control_flow_ops.Assert(i < 10, [i, [10], [i + 1]])
def kl_divergence(distribution_a,
distribution_b,
allow_nan_stats=True,
name=None):
"""Get the KL-divergence KL(distribution_a || distribution_b).
If there is no KL method registered specifically for `type(distribution_a)`
and `type(distribution_b)`, then the class hierarchies of these types are
searched.
If one KL method is registered between any pairs of classes in these two
parent hierarchies, it is used.
If more than one such registered method exists, the method whose registered
classes have the shortest sum MRO paths to the input types is used.
If more than one such shortest path exists, the first method
identified in the search is used (favoring a shorter MRO distance to
`type(distribution_a)`).
Args:
distribution_a: The first distribution.
distribution_b: The second distribution.
allow_nan_stats: Python `bool`, default `True`. When `True`,
statistics (e.g., mean, mode, variance) use the value "`NaN`" to
indicate the result is undefined. When `False`, an exception is raised
if one or more of the statistic's batch members are undefined.
name: Python `str` name prefixed to Ops created by this class.
Returns:
A Tensor with the batchwise KL-divergence between `distribution_a`
and `distribution_b`.
Raises:
NotImplementedError: If no KL method is defined for distribution types
of `distribution_a` and `distribution_b`.
"""
kl_fn = _registered_kl(type(distribution_a), type(distribution_b))
if kl_fn is None:
raise NotImplementedError(
"No KL(distribution_a || distribution_b) registered for distribution_a "
"type %s and distribution_b type %s" %
(type(distribution_a).__name__, type(distribution_b).__name__))
with ops.name_scope("KullbackLeibler"):
kl_t = kl_fn(distribution_a, distribution_b, name=name)
if allow_nan_stats:
return kl_t
# Check KL for NaNs
kl_t = array_ops.identity(kl_t, name="kl")
with ops.control_dependencies([
control_flow_ops.Assert(
math_ops.logical_not(
math_ops.reduce_any(math_ops.is_nan(kl_t))),
[
"KL calculation between %s and %s returned NaN values "
"(and was called with allow_nan_stats=False). Values:"
% (distribution_a.name, distribution_b.name), kl_t
])
]):
return array_ops.identity(kl_t, name="checked_kl")
def _training_examples_and_variables():
"""Returns dictionaries for training examples and variables."""
batch_size = targets.get_shape()[0]
# Iterate over all feature columns and create appropriate lists for dense
# and sparse features as well as dense and sparse weights (variables) for
# SDCA.
# TODO(sibyl-vie3Poto): Reshape variables stored as values in column_to_variables
# dict as 1-dimensional tensors.
dense_features, sparse_features, sparse_feature_with_values = [], [], []
dense_feature_weights = []
sparse_feature_weights, sparse_feature_with_values_weights = [], []
for column in sorted(columns_to_variables.keys(), key=lambda x: x.key):
transformed_tensor = features[column]
if isinstance(column, layers.feature_column._RealValuedColumn): # pylint: disable=protected-access
# A real-valued column corresponds to a dense feature in SDCA. A
# transformed tensor corresponding to a RealValuedColumn should have
# rank at most 2. In order to be passed to SDCA, its rank needs to be
# exactly 2 (i.e., its shape should be [batch_size, column.dim]).
check_rank_op = control_flow_ops.Assert(
math_ops.less_equal(array_ops.rank(transformed_tensor), 2),
['transformed_tensor shouls have rank at most 2.'])
# Reshape to [batch_size, dense_column_dimension].
with ops.control_dependencies([check_rank_op]):
transformed_tensor = array_ops.reshape(transformed_tensor, [
array_ops.shape(transformed_tensor)[0], -1
])
dense_features.append(transformed_tensor)
# For real valued columns, the variables list contains exactly one
# element.
dense_feature_weights.append(columns_to_variables[column][0])
elif isinstance(column, layers.feature_column._BucketizedColumn): # pylint: disable=protected-access
# A bucketized column corresponds to a sparse feature in SDCA. The
# bucketized feature is "sparsified" for SDCA by converting it to a
# SparseFeatureColumn respresenting the one-hot encoding of the
# bucketized feature.
#
# TODO(sibyl-vie3Poto): Explore whether it is more efficient to translate a
# bucketized feature column to a dense feature in SDCA. This will
# likely depend on the number of buckets.
dense_bucket_tensor = column._to_dnn_input_layer(transformed_tensor) # pylint: disable=protected-access
sparse_feature_column = _dense_tensor_to_sparse_feature_column(
dense_bucket_tensor)
sparse_feature_with_values.append(sparse_feature_column)
# For bucketized columns, the variables list contains exactly one
# element.
sparse_feature_with_values_weights.append(
columns_to_variables[column][0])
elif isinstance(
column,
(
layers.feature_column._CrossedColumn, # pylint: disable=protected-access
layers.feature_column._SparseColumn)): # pylint: disable=protected-access
sparse_features.append(
SparseFeatureColumn(
array_ops.reshape(
array_ops.split(
value=transformed_tensor.indices,
num_or_size_splits=2,
axis=1)[0], [-1]),
array_ops.reshape(transformed_tensor.values, [-1]), None))
sparse_feature_weights.append(columns_to_variables[column][0])
elif isinstance(column, layers.feature_column._WeightedSparseColumn): # pylint: disable=protected-access
id_tensor = column.id_tensor(transformed_tensor)
weight_tensor = column.weight_tensor(transformed_tensor)
sparse_feature_with_values.append(
SparseFeatureColumn(
array_ops.reshape(
array_ops.split(
value=id_tensor.indices, num_or_size_splits=2, axis=1)
[0], [-1]),
array_ops.reshape(id_tensor.values, [-1]),
array_ops.reshape(weight_tensor.values, [-1])))
sparse_feature_with_values_weights.append(
columns_to_variables[column][0])
else:
raise ValueError('SDCAOptimizer does not support column type %s.' %
type(column).__name__)
example_weights = array_ops.reshape(
features[weight_column_name],
shape=[-1]) if weight_column_name else array_ops.ones([batch_size])
example_ids = features[self._example_id_column]
sparse_feature_with_values.extend(sparse_features)
sparse_feature_with_values_weights.extend(sparse_feature_weights)
examples = dict(
sparse_features=sparse_feature_with_values,
dense_features=dense_features,
example_labels=math_ops.to_float(
array_ops.reshape(targets, shape=[-1])),
example_weights=example_weights,
example_ids=example_ids)
sdca_variables = dict(
sparse_features_weights=sparse_feature_with_values_weights,
dense_features_weights=dense_feature_weights)
return examples, sdca_variables
def broadcast_dimension(self, axis, lengths):
"""Returns a shape that is broadcast-compatible with self & lengths.
* If dimension[axis] is uniform and lengths is a scalar, the check
that either lengths==1 or axis==1 or lengths==axis, and tile
dimension[axis] with tf.where(lengths==axis, 1, axis) repeats.
* If dimension[axis] is uniform and lengths is a vector, then check
that dimension[axis]==1, and raggedly tile dimension[axis] with
lengths repeats. (we can skip tiling if we statically know that
slice_lengths == 1??)
* If dimension[axis] is ragged and lengths is a scalar, then check
that lengths==1.
* If dimension[axis] is ragged and lengths is a vector, then check
that self.dimension_size(axis) == lengths.
Args:
axis: `int`. The dimension to broadcast.
lengths: 0-D or 1-D integer `Tensor`.
Returns:
A `RaggedTensorDynamicShape`.
"""
lengths = ragged_util.convert_to_int_tensor(lengths,
name='lengths',
dtype=self.dim_size_dtype)
# Check whether lengths is a scalar (for uniform dimensions) or
# vector (for ragged dimensions).
if lengths.shape.ndims is None:
raise ValueError('lengths must have a known rank.')
elif lengths.shape.ndims > 1:
raise ValueError('lengths must be a scalar or vector')
else:
lengths_is_scalar = (lengths.shape.ndims == 0)
# Verify that the shapes are compatible.
if self.is_ragged(axis):
if lengths_is_scalar:
condition = math_ops.equal(lengths, 1)
else:
condition = math_ops.reduce_all(
math_ops.equal(lengths, self.dimension_size(axis)))
else:
axis_dim_size = self.dimension_size(axis)
if lengths_is_scalar:
condition = (math_ops.equal(lengths, 1)
| math_ops.equal(axis_dim_size, 1)
| math_ops.equal(axis_dim_size, lengths))
else:
condition = math_ops.equal(axis_dim_size, 1)
broadcast_err = [
'Unable to broadcast: dimension size mismatch in dimension', axis,
'lengths=', lengths, 'dim_size=',
self.dimension_size(axis)
]
broadcast_check = control_flow_ops.Assert(condition,
data=broadcast_err,
summarize=10)
with ops.control_dependencies([broadcast_check]):
# Partitioned dimensions:
if axis < self.num_partitioned_dimensions:
if self.is_ragged(axis):
# Use an identity op to make sure the check actually gets run.
return RaggedTensorDynamicShape(
self._partitioned_dim_sizes,
array_ops.identity(self.inner_dim_sizes))
else:
return self._broadcast_uniform_partitioned_dimension(
axis, lengths)
# Inner dimensions:
else:
if lengths_is_scalar:
return self._broadcast_inner_dimension_to_uniform(
axis, lengths)
else:
if axis == 0:
raise ValueError(
'Unable to broadcast: '
'outermost dimension must be uniform.')
return self._broadcast_inner_dimension_to_ragged(
axis, lengths)
def random_with_assert(x):
y = random(x)
assert_op = control_flow_ops.Assert(math_ops.greater_equal(y, 1), [y])
with ops.control_dependencies([assert_op]):
return y
def MyFn(x):
with ops.control_dependencies(
[control_flow_ops.Assert(math_ops.less_equal(x, 10.0), [x])]):
return array_ops.identity(x)
def func():
control_flow_ops.Assert(False, [False])
return
def body(it):
nonlocal i
with ops.control_dependencies(
(control_flow_ops.Assert(i < 3, (i, )), )):
i = it
def fail(i):
control_flow_ops.Assert(math_ops.equal(i, 0), ['ick'])
def _get_update_op(self,
score_drop,
score_grow,
mask,
weights,
reinit_when_same=False):
"""Prunes+grows connections, all tensors same shape."""
old_dtype = mask.dtype
mask_casted = math_ops.cast(mask, dtypes.float32)
n_total = array_ops.size(score_drop)
n_ones = math_ops.cast(math_ops.reduce_sum(mask_casted),
dtype=dtypes.int32)
n_prune = math_ops.cast(
math_ops.cast(n_ones, dtype=dtypes.float32) * self.drop_fraction,
dtypes.int32)
n_keep = n_ones - n_prune
# Sort the entire array since the k needs to be constant for TPU.
_, sorted_indices = nn_ops.top_k(array_ops.reshape(score_drop, [-1]),
k=n_total)
sorted_indices_ex = array_ops.expand_dims(sorted_indices, 1)
# We will have zeros after having `n_keep` many ones.
new_values = array_ops.where(
math_ops.range(n_total) < n_keep,
array_ops.ones_like(sorted_indices, dtype=mask_casted.dtype),
array_ops.zeros_like(sorted_indices, dtype=mask_casted.dtype))
mask1 = array_ops.scatter_nd(sorted_indices_ex, new_values,
new_values.shape)
# Flatten the scores
score_grow = array_ops.reshape(score_grow, [-1])
# Set scores of the enabled connections(ones) to min(s) - 1, so that they
# have the lowest scores.
score_grow_lifted = array_ops.where(
math_ops.equal(mask1, 1),
array_ops.ones_like(mask1) * (math_ops.reduce_min(score_grow) - 1),
score_grow)
_, sorted_indices = nn_ops.top_k(score_grow_lifted, k=n_total)
sorted_indices_ex = array_ops.expand_dims(sorted_indices, 1)
new_values = array_ops.where(
math_ops.range(n_total) < n_prune,
array_ops.ones_like(sorted_indices, dtype=mask_casted.dtype),
array_ops.zeros_like(sorted_indices, dtype=mask_casted.dtype))
mask2 = array_ops.scatter_nd(sorted_indices_ex, new_values,
new_values.shape)
# Ensure masks are disjoint.
assert_op = control_flow_ops.Assert(
math_ops.equal(math_ops.reduce_sum(mask1 * mask2), 0.),
[mask1, mask2])
with ops.control_dependencies([assert_op]):
# Let's set the weights of the growed connections.
mask2_reshaped = array_ops.reshape(mask2, mask.shape)
# Set the values of the new connections.
grow_tensor = self.get_grow_tensor(weights, self._grow_init)
if reinit_when_same:
# If dropped and grown, we re-initialize.
new_connections = math_ops.equal(mask2_reshaped, 1)
else:
new_connections = math_ops.logical_and(
math_ops.equal(mask2_reshaped, 1),
math_ops.equal(mask_casted, 0))
new_weights = array_ops.where(new_connections, grow_tensor, weights)
weights_update = state_ops.assign(weights, new_weights)
# Ensure there is no momentum value for new connections
reset_op = self.reset_momentum(weights, new_connections)
with ops.control_dependencies([weights_update, reset_op]):
mask_combined = array_ops.reshape(mask1 + mask2, mask.shape)
mask_combined = math_ops.cast(mask_combined, dtype=old_dtype)
new_mask = state_ops.assign(mask, mask_combined)
return new_mask
def _assert_increasing(t):
assert_increasing = control_flow_ops.Assert(
math_ops.reduce_all(t[1:] > t[:-1]),
['`t` must be monotonic increasing'])
return ops.control_dependencies([assert_increasing])
def raise_zero_rank_error():
msg = gen_string_ops.string_join(
['len requires non-zero rank, got ',
gen_string_ops.as_string(rank)])
with ops.control_dependencies([control_flow_ops.Assert(False, [msg])]):
return constant_op.constant(0, dtype=dtypes.int32)
