def _inverse_log_det_jacobian(self, y, **kwargs):
ildj = constant_op.constant(
0., dtype=y.dtype.base_dtype, name="inverse_log_det_jacobian")
if not self.bijectors:
return ildj
event_ndims = _maybe_get_event_ndims_statically(
self.inverse_min_event_ndims)
if _use_static_shape(y, event_ndims):
event_shape = y.shape[y.shape.ndims - event_ndims:]
else:
event_shape = array_ops.shape(y)[array_ops.rank(y) - event_ndims:]
for b in self.bijectors:
ildj += b.inverse_log_det_jacobian(
y, event_ndims=event_ndims, **kwargs.get(b.name, {}))
if _use_static_shape(y, event_ndims):
event_shape = b.inverse_event_shape(event_shape)
event_ndims = _maybe_get_event_ndims_statically(event_shape.ndims)
else:
event_shape = b.inverse_event_shape_tensor(event_shape)
event_ndims = _maybe_get_event_ndims_statically(
array_ops.rank(event_shape))
y = b.inverse(y, **kwargs.get(b.name, {}))
return ildj
def same_dynamic_shape(a, b):
"""Returns whether a and b have the same dynamic shape.
Args:
a: `Tensor`
b: `Tensor`
Returns:
`Boolean` `Tensor` representing if both tensors have the same shape.
"""
a = ops.convert_to_tensor(a, name="a")
b = ops.convert_to_tensor(b, name="b")
# Here we can't just do math_ops.equal(a.shape, b.shape), since
# static shape inference may break the equality comparison between
# shape(a) and shape(b) in math_ops.equal.
def all_shapes_equal():
return math_ops.reduce_all(math_ops.equal(
array_ops.concat([array_ops.shape(a), array_ops.shape(b)], 0),
array_ops.concat([array_ops.shape(b), array_ops.shape(a)], 0)))
# One of the shapes isn't fully defined, so we need to use the dynamic
# shape.
return control_flow_ops.cond(
math_ops.equal(array_ops.rank(a), array_ops.rank(b)),
all_shapes_equal,
lambda: constant_op.constant(False))
def _forward_log_det_jacobian(self, x, **kwargs):
x = ops.convert_to_tensor(x, name="x")
fldj = constant_op.constant(
0., dtype=x.dtype, name="inverse_log_det_jacobian")
if not self.bijectors:
return fldj
event_ndims = _maybe_get_event_ndims_statically(
self.forward_min_event_ndims)
if _use_static_shape(x, event_ndims):
event_shape = x.shape[x.shape.ndims - event_ndims:]
else:
event_shape = array_ops.shape(x)[array_ops.rank(x) - event_ndims:]
for b in reversed(self.bijectors):
fldj += b.forward_log_det_jacobian(
x, event_ndims=event_ndims, **kwargs.get(b.name, {}))
if _use_static_shape(x, event_ndims):
event_shape = b.forward_event_shape(event_shape)
event_ndims = _maybe_get_event_ndims_statically(event_shape.ndims)
else:
event_shape = b.forward_event_shape_tensor(event_shape)
event_ndims = _maybe_get_event_ndims_statically(
array_ops.rank(event_shape))
x = b.forward(x, **kwargs.get(b.name, {}))
return fldj
def rank(input, name=None): # pylint: disable=redefined-builtin
"""Returns the rank of a RaggedTensor.
Returns a 0-D `int32` `Tensor` representing the rank of `input`.
For example:
```python
# shape of tensor 't' is [2, None, None]
t = tf.ragged.constant([[[1], [2, 2]], [[3, 3, 3], [4, 4, 4, 4]]])
tf.rank(t) # 3
```
Args:
input: A `RaggedTensor`
name: A name for the operation (optional).
Returns:
A `Tensor` of type `int32`.
"""
with ops.name_scope(name, 'RaggedRank', [input]) as name:
if not ragged_tensor.is_ragged(input):
return array_ops.rank(input, name)
return input.ragged_rank + array_ops.rank(input.flat_values)
def _check_labels_and_scores(boolean_labels, scores, check_shape):
"""Check the rank of labels/scores, return tensor versions."""
with ops.op_scope([boolean_labels, scores], '_check_labels_and_scores'):
boolean_labels = ops.convert_to_tensor(boolean_labels,
name='boolean_labels')
scores = ops.convert_to_tensor(scores, name='scores')
if boolean_labels.dtype != dtypes.bool:
raise ValueError(
'Argument boolean_labels should have dtype bool. Found: %s',
boolean_labels.dtype)
if check_shape:
labels_rank_1 = logging_ops.Assert(
math_ops.equal(1, array_ops.rank(boolean_labels)),
['Argument boolean_labels should have rank 1. Found: ',
boolean_labels.name, array_ops.shape(boolean_labels)])
scores_rank_1 = logging_ops.Assert(
math_ops.equal(1, array_ops.rank(scores)),
['Argument scores should have rank 1. Found: ', scores.name,
array_ops.shape(scores)])
with ops.control_dependencies([labels_rank_1, scores_rank_1]):
return boolean_labels, scores
else:
return boolean_labels, scores
def _check_mu(self, mu):
"""Return `mu` after validity checks and possibly with assertations."""
mu = ops.convert_to_tensor(mu)
cov = self._cov
if mu.dtype != cov.dtype:
raise TypeError(
"mu and cov must have the same dtype. Found mu.dtype = %s, "
"cov.dtype = %s"
% (mu.dtype, cov.dtype))
if not self.strict:
return mu
else:
assert_compatible_shapes = control_flow_ops.group(
check_ops.assert_equal(
array_ops.rank(mu) + 1,
cov.rank(),
data=["mu should have rank 1 less than cov. Found: rank(mu) = ",
array_ops.rank(mu), " rank(cov) = ", cov.rank()],
),
check_ops.assert_equal(
array_ops.shape(mu),
cov.vector_shape(),
data=["mu.shape and cov.shape[:-1] should match. "
"Found: shape(mu) = "
, array_ops.shape(mu), " shape(cov) = ", cov.shape()],
),
)
return control_flow_ops.with_dependencies([assert_compatible_shapes], mu)
def _get_sparse_tensors(self, inputs, weight_collections=None,
trainable=None):
sparse_tensors = self.categorical_column._get_sparse_tensors(inputs)
id_tensor = sparse_tensors.id_tensor
weight_tensor = sparse_tensors.weight_tensor
# Expands final dimension, so that embeddings are not combined during
# embedding lookup.
check_id_rank = check_ops.assert_equal(
array_ops.rank(id_tensor), 2,
data=[
'Column {} expected ID tensor of rank 2. '.format(self.name),
'id_tensor shape: ', array_ops.shape(id_tensor)])
with ops.control_dependencies([check_id_rank]):
id_tensor = sparse_ops.sparse_reshape(
id_tensor,
shape=array_ops.concat([id_tensor.dense_shape, [1]], axis=0))
if weight_tensor is not None:
check_weight_rank = check_ops.assert_equal(
array_ops.rank(weight_tensor), 2,
data=[
'Column {} expected weight tensor of rank 2.'.format(self.name),
'weight_tensor shape:', array_ops.shape(weight_tensor)])
with ops.control_dependencies([check_weight_rank]):
weight_tensor = sparse_ops.sparse_reshape(
weight_tensor,
shape=array_ops.concat([weight_tensor.dense_shape, [1]], axis=0))
return fc._CategoricalColumn.IdWeightPair(id_tensor, weight_tensor)
def _check_shapes_dynamic(self, operator, v, diag):
"""Return (v, diag) with Assert dependencies, which check shape."""
checks = []
with ops.op_scope([operator, v, diag], 'check_shapes'):
s_v = array_ops.shape(v)
r_op = operator.rank()
r_v = array_ops.rank(v)
if diag is not None:
s_d = array_ops.shape(diag)
r_d = array_ops.rank(diag)
# Check tensor rank.
checks.append(check_ops.assert_rank(v, r_op))
if diag is not None:
checks.append(check_ops.assert_rank(diag, r_op - 1))
# Check batch shape
checks.append(check_ops.assert_equal(
operator.batch_shape(), array_ops.slice(s_v, [0], [r_v - 2])))
if diag is not None:
checks.append(check_ops.assert_equal(
operator.batch_shape(), array_ops.slice(s_d, [0], [r_d - 1])))
# Check event shape
checks.append(check_ops.assert_equal(
operator.vector_space_dimension(), array_ops.gather(s_v, r_v - 2)))
if diag is not None:
checks.append(check_ops.assert_equal(
array_ops.gather(s_v, r_v - 1), array_ops.gather(s_d, r_d - 1)))
v = control_flow_ops.with_dependencies(checks, v)
if diag is not None:
diag = control_flow_ops.with_dependencies(checks, diag)
return v, diag
def check(t): target = array_ops.shape(tensor)[1:] result = array_ops.broadcast_dynamic_shape(target, array_ops.shape(t)) # This rank check ensures that I don't get a wrong answer from the # _shapes_ broadcasting against each other. gt = check_ops.assert_greater(array_ops.rank(target), array_ops.rank(t)) eq = check_ops.assert_equal(target, result) return gt, eq
def sign_magnitude_positive_definite(
raw, off_diagonal_scale=0., overall_scale=0.):
"""Constructs a positive definite matrix from an unconstrained input matrix.
We want to keep the whole matrix on a log scale, but also allow off-diagonal
elements to be negative, so the sign of off-diagonal elements is modeled
separately from their magnitude (using the lower and upper triangles
respectively). Specifically:
for i < j, we have:
output_cholesky[i, j] = raw[j, i] / (abs(raw[j, i]) + 1) *
exp((off_diagonal_scale + overall_scale + raw[i, j]) / 2)
output_cholesky[i, i] = exp((raw[i, i] + overall_scale) / 2)
output = output_cholesky^T * output_cholesky
where raw, off_diagonal_scale, and overall_scale are
un-constrained real-valued variables. The resulting values are stable
around zero due to the exponential (and the softsign keeps the function
smooth).
Args:
raw: A [..., M, M] Tensor.
off_diagonal_scale: A scalar or [...] shaped Tensor controlling the relative
scale of off-diagonal values in the output matrix.
overall_scale: A scalar or [...] shaped Tensor controlling the overall scale
of the output matrix.
Returns:
The `output` matrix described above, a [..., M, M] positive definite matrix.
"""
raw = ops.convert_to_tensor(raw)
diagonal = array_ops.matrix_diag_part(raw)
def _right_pad_with_ones(tensor, target_rank):
# Allow broadcasting even if overall_scale and off_diagonal_scale have batch
# dimensions
tensor = ops.convert_to_tensor(tensor, dtype=raw.dtype.base_dtype)
return array_ops.reshape(tensor,
array_ops.concat(
[
array_ops.shape(tensor), array_ops.ones(
[target_rank - array_ops.rank(tensor)],
dtype=target_rank.dtype)
],
axis=0))
# We divide the log values by 2 to compensate for the squaring that happens
# when transforming Cholesky factors into positive definite matrices.
sign_magnitude = (gen_math_ops.exp(
(raw + _right_pad_with_ones(off_diagonal_scale, array_ops.rank(raw)) +
_right_pad_with_ones(overall_scale, array_ops.rank(raw))) / 2.) *
nn.softsign(array_ops.matrix_transpose(raw)))
sign_magnitude.set_shape(raw.get_shape())
cholesky_factor = array_ops.matrix_set_diag(
input=array_ops.matrix_band_part(sign_magnitude, 0, -1),
diagonal=gen_math_ops.exp((diagonal + _right_pad_with_ones(
overall_scale, array_ops.rank(diagonal))) / 2.))
return math_ops.matmul(cholesky_factor, cholesky_factor, transpose_a=True)
def remove_squeezable_dimensions(
labels, predictions, expected_rank_diff=0, name=None):
"""Squeeze last dim if ranks differ from expected by exactly 1.
In the common case where we expect shapes to match, `expected_rank_diff`
defaults to 0, and we squeeze the last dimension of the larger rank if they
differ by 1.
But, for example, if `labels` contains class IDs and `predictions` contains 1
probability per class, we expect `predictions` to have 1 more dimension than
`labels`, so `expected_rank_diff` would be 1. In this case, we'd squeeze
`labels` if `rank(predictions) - rank(labels) == 0`, and
`predictions` if `rank(predictions) - rank(labels) == 2`.
This will use static shape if available. Otherwise, it will add graph
operations, which could result in a performance hit.
Args:
labels: Label values, a `Tensor` whose dimensions match `predictions`.
predictions: Predicted values, a `Tensor` of arbitrary dimensions.
expected_rank_diff: Expected result of `rank(predictions) - rank(labels)`.
name: Name of the op.
Returns:
Tuple of `labels` and `predictions`, possibly with last dim squeezed.
"""
with ops.name_scope(name, 'remove_squeezable_dimensions',
[labels, predictions]):
predictions = ops.convert_to_tensor(predictions)
labels = ops.convert_to_tensor(labels)
predictions_shape = predictions.get_shape()
predictions_rank = predictions_shape.ndims
labels_shape = labels.get_shape()
labels_rank = labels_shape.ndims
if (labels_rank is not None) and (predictions_rank is not None):
# Use static rank.
rank_diff = predictions_rank - labels_rank
if rank_diff == expected_rank_diff + 1:
predictions = array_ops.squeeze(predictions, [-1])
elif rank_diff == expected_rank_diff - 1:
labels = array_ops.squeeze(labels, [-1])
return labels, predictions
# Use dynamic rank.
rank_diff = array_ops.rank(predictions) - array_ops.rank(labels)
if (predictions_rank is None) or (
predictions_shape.dims[-1].is_compatible_with(1)):
predictions = control_flow_ops.cond(
math_ops.equal(expected_rank_diff + 1, rank_diff),
lambda: array_ops.squeeze(predictions, [-1]),
lambda: predictions)
if (labels_rank is None) or (
labels_shape.dims[-1].is_compatible_with(1)):
labels = control_flow_ops.cond(
math_ops.equal(expected_rank_diff - 1, rank_diff),
lambda: array_ops.squeeze(labels, [-1]),
lambda: labels)
return labels, predictions
def testRank(self):
with self.test_scope():
self.assertEqual(
0, array_ops.rank(constant_op.constant(1.0)).numpy())
self.assertEqual(
1, array_ops.rank(constant_op.constant([1.0, 2.0, 3.0])).numpy())
self.assertEqual(
2, array_ops.rank(
constant_op.constant([[1.0, 2.0], [3.0, 4.0]])).numpy())
def _get_chol_and_x_compatible_shape(self, x):
"""Return self.chol and x, (possibly) broadcast to compatible shape."""
# x and chol are "compatible" if their shape matches except for the last two
# dimensions of chol are [k, k], and the last two of x are [k, 1].
# E.g. x.shape = [A, B, k, 1], and chol.shape = [A, B, k, k]
# This is required for the batch_triangular_solve, which does not broadcast.
# TODO(langmore) This broadcast replicates matrices unnecesarily! In the
# case where
# x.shape = [M1,...,Mr, N1,...,Nb, k], and chol.shape = [N1,...,Nb, k, k]
# (which is common if x was sampled), the front dimensions of x can be
# "flipped" to the end, making
# x_flipped.shape = [N1,...,Nb, k, M1*...*Mr],
# and this can be handled by the linear solvers. This is preferred, because
# it does not replicate the matrix, or create any new data.
# We assume x starts without the trailing singleton dimension, e.g.
# x.shape = [B, k].
chol = self._chol
with ops.op_scope([x] + self.inputs, 'get_chol_and_x_compatible_shape'):
# If we determine statically that shapes match, we're done.
if x.get_shape() == chol.get_shape()[:-1]:
x_expanded = array_ops.expand_dims(x, -1)
return chol, x_expanded
# Dynamic check if shapes match or not.
vector_shape = self.vector_shape() # Shape of chol minus last dim.
are_same_rank = math_ops.equal(
array_ops.rank(x), array_ops.rank(vector_shape))
def shapes_match_if_same_rank():
return math_ops.reduce_all(math_ops.equal(
array_ops.shape(x), vector_shape))
shapes_match = control_flow_ops.cond(are_same_rank,
shapes_match_if_same_rank,
lambda: ops.convert_to_tensor(False))
# Make tensors (never instantiated) holding the broadcast shape.
# matrix_broadcast_dummy is the shape we will broadcast chol to.
matrix_bcast_dummy = chol + array_ops.expand_dims(x, -1)
# vector_bcast_dummy is the shape we will bcast x to, before we expand it.
chol_minus_last_dim = math_ops.reduce_sum(chol, reduction_indices=[-1])
vector_bcast_dummy = x + chol_minus_last_dim
chol_bcast = chol + array_ops.zeros_like(matrix_bcast_dummy)
x_bcast = x + array_ops.zeros_like(vector_bcast_dummy)
chol_result = control_flow_ops.cond(shapes_match, lambda: chol,
lambda: chol_bcast)
chol_result.set_shape(matrix_bcast_dummy.get_shape())
x_result = control_flow_ops.cond(shapes_match, lambda: x, lambda: x_bcast)
x_result.set_shape(vector_bcast_dummy.get_shape())
x_expanded = array_ops.expand_dims(x_result, -1)
return chol_result, x_expanded
def _reshape_helper(self, x, event_shape_in, event_shape_out):
"""Reshape only the event_shape of an input `Tensor`."""
def _get_rank_from_shape(shape):
"""Computes rank from a shape `Tensor`, statically if possible."""
# Uses fact that rank is "shape of shape".
ndims = shape.shape.with_rank_at_least(1)[0].value
if ndims is not None:
return ndims, ndims
return None, array_ops.shape(shape)[0]
event_ndims_in_, event_ndims_in = _get_rank_from_shape(event_shape_in)
assertions = []
# Ensure x.event_shape is compatible with event_shape_in.
if x.shape.ndims is not None:
x_ndims_, x_ndims = [x.shape.ndims]*2
else:
x_ndims_, x_ndims = None, array_ops.rank(x)
if (event_ndims_in_ is not None
and x_ndims_ is not None
and x.shape.with_rank_at_least(event_ndims_in_)[
x_ndims_-event_ndims_in_:].is_fully_defined()):
x_event_shape_, x_event_shape = [ # pylint: disable=unbalanced-tuple-unpacking
np.int32(x.shape[x_ndims_-event_ndims_in_:])]*2
else:
x_event_shape_, x_event_shape = (
None, array_ops.shape(x)[x_ndims-event_ndims_in:])
event_shape_in_ = tensor_util.constant_value(event_shape_in)
if x_event_shape_ is not None and event_shape_in_ is not None:
if not np.equal(x_event_shape_, event_shape_in_).all():
raise ValueError(
"Input `event_shape` ({}) does not match `event_shape_in` ({}).".
format(x_event_shape_, event_shape_in_))
elif self.validate_args:
assertions.append(check_ops.assert_equal(
x_event_shape, event_shape_in,
message="Input `event_shape` does not match `event_shape_in`."))
if assertions:
x = control_flow_ops.with_dependencies(assertions, x)
# get the parts of shape(x) that will not change
sample_and_batch_shape = array_ops.shape(x)
ndims = (x.shape.ndims if x.shape.ndims is not None
else array_ops.rank(x))
sample_and_batch_shape = sample_and_batch_shape[
:(ndims - math_ops.abs(event_ndims_in))]
new_shape = array_ops.concat(
[sample_and_batch_shape, event_shape_out], axis=0)
return array_ops.reshape(x, new_shape)
def testDenseShape(self):
t_value = [[0, 42], [24, 0]]
self.assertAllEqual((2, 2), self.evaluate(array_ops.shape(t_value)))
self.assertEqual(4, self.evaluate(array_ops.size(t_value)))
self.assertEqual(2, self.evaluate(array_ops.rank(t_value)))
t = constant_op.constant(t_value)
self.assertAllEqual((2, 2), self.evaluate(array_ops.shape(t)))
self.assertEqual(4, self.evaluate(array_ops.size(t)))
self.assertEqual(2, self.evaluate(array_ops.rank(t)))
def testDenseShape(self):
with self.test_session():
t_value = [[0, 42], [24, 0]]
self.assertAllEqual((2, 2), array_ops.shape(t_value).eval())
self.assertEqual(4, array_ops.size(t_value).eval())
self.assertEqual(2, array_ops.rank(t_value).eval())
t = constant_op.constant(t_value)
self.assertAllEqual((2, 2), array_ops.shape(t).eval())
self.assertEqual(4, array_ops.size(t).eval())
self.assertEqual(2, array_ops.rank(t).eval())
def testSparseShape(self):
sp_value = sparse_tensor.SparseTensorValue(
indices=((0, 1), (1, 0)), values=(42, 24), dense_shape=(2, 2))
self.assertAllEqual((2, 2), self.evaluate(array_ops.shape(sp_value)))
self.assertEqual(4, self.evaluate(array_ops.size(sp_value)))
self.assertEqual(2, self.evaluate(array_ops.rank(sp_value)))
sp = sparse_tensor.SparseTensor.from_value(sp_value)
self.assertAllEqual((2, 2), self.evaluate(array_ops.shape(sp)))
self.assertEqual(4, self.evaluate(array_ops.size(sp)))
self.assertEqual(2, self.evaluate(array_ops.rank(sp)))
def testSparseShape(self):
with self.test_session():
sp_value = sparse_tensor.SparseTensorValue(
indices=((0, 1), (1, 0)), values=(42, 24), dense_shape=(2, 2))
self.assertAllEqual((2, 2), array_ops.shape(sp_value).eval())
self.assertEqual(4, array_ops.size(sp_value).eval())
self.assertEqual(2, array_ops.rank(sp_value).eval())
sp = sparse_tensor.SparseTensor.from_value(sp_value)
self.assertAllEqual((2, 2), array_ops.shape(sp).eval())
self.assertEqual(4, array_ops.size(sp).eval())
self.assertEqual(2, array_ops.rank(sp).eval())
def _compute_energy_change(current_target_log_prob,
current_momentums,
proposed_target_log_prob,
proposed_momentums,
independent_chain_ndims,
name=None):
"""Helper to `kernel` which computes the energy change."""
with ops.name_scope(
name, "compute_energy_change",
([current_target_log_prob, proposed_target_log_prob,
independent_chain_ndims] +
current_momentums + proposed_momentums)):
# Abbreviate lk0=log_kinetic_energy and lk1=proposed_log_kinetic_energy
# since they're a mouthful and lets us inline more.
lk0, lk1 = [], []
for current_momentum, proposed_momentum in zip(current_momentums,
proposed_momentums):
axis = math_ops.range(independent_chain_ndims,
array_ops.rank(current_momentum))
lk0.append(_log_sum_sq(current_momentum, axis))
lk1.append(_log_sum_sq(proposed_momentum, axis))
lk0 = -np.log(2.) + math_ops.reduce_logsumexp(array_ops.stack(lk0, axis=-1),
axis=-1)
lk1 = -np.log(2.) + math_ops.reduce_logsumexp(array_ops.stack(lk1, axis=-1),
axis=-1)
lp0 = -current_target_log_prob # log_potential
lp1 = -proposed_target_log_prob # proposed_log_potential
x = array_ops.stack([lp1, math_ops.exp(lk1), -lp0, -math_ops.exp(lk0)],
axis=-1)
# The sum is NaN if any element is NaN or we see both +Inf and -Inf.
# Thus we will replace such rows with infinite energy change which implies
# rejection. Recall that float-comparisons with NaN are always False.
is_sum_determinate = (
math_ops.reduce_all(math_ops.is_finite(x) | (x >= 0.), axis=-1) &
math_ops.reduce_all(math_ops.is_finite(x) | (x <= 0.), axis=-1))
is_sum_determinate = array_ops.tile(
is_sum_determinate[..., array_ops.newaxis],
multiples=array_ops.concat([
array_ops.ones(array_ops.rank(is_sum_determinate),
dtype=dtypes.int32),
[4],
], axis=0))
x = array_ops.where(is_sum_determinate,
x,
array_ops.fill(array_ops.shape(x),
value=x.dtype.as_numpy_dtype(np.inf)))
return math_ops.reduce_sum(x, axis=-1)
def run_inception(images,
graph_def=None,
default_graph_def_fn=_default_graph_def_fn,
image_size=INCEPTION_DEFAULT_IMAGE_SIZE,
input_tensor=INCEPTION_INPUT,
output_tensor=INCEPTION_OUTPUT):
"""Run images through a pretrained Inception classifier.
Args:
images: Input tensors. Must be [batch, height, width, channels]. Input shape
and values must be in [-1, 1], which can be achieved using
`preprocess_image`.
graph_def: A GraphDef proto of a pretrained Inception graph. If `None`,
call `default_graph_def_fn` to get GraphDef.
default_graph_def_fn: A function that returns a GraphDef. Used if
`graph_def` is `None. By default, returns a pretrained InceptionV3 graph.
image_size: Required image width and height. See unit tests for the default
values.
input_tensor: Name of input Tensor.
output_tensor: Name or list of output Tensors. This function will compute
activations at the specified layer. Examples include INCEPTION_V3_OUTPUT
and INCEPTION_V3_FINAL_POOL which would result in this function computing
the final logits or the penultimate pooling layer.
Returns:
Tensor or Tensors corresponding to computed `output_tensor`.
Raises:
ValueError: If images are not the correct size.
ValueError: If neither `graph_def` nor `default_graph_def_fn` are provided.
"""
images = _validate_images(images, image_size)
if graph_def is None:
if default_graph_def_fn is None:
raise ValueError('If `graph_def` is `None`, must provide '
'`default_graph_def_fn`.')
graph_def = default_graph_def_fn()
activations = run_image_classifier(images, graph_def, input_tensor,
output_tensor)
if isinstance(activations, list):
for i, activation in enumerate(activations):
if array_ops.rank(activation) != 2:
activations[i] = layers.flatten(activation)
else:
if array_ops.rank(activations) != 2:
activations = layers.flatten(activations)
return activations
def _expand_is_accepted_like(x):
with ops.name_scope("_choose"):
expand_shape = array_ops.concat([
array_ops.shape(is_accepted),
array_ops.ones([array_ops.rank(x) - array_ops.rank(is_accepted)],
dtype=dtypes.int32),
], axis=0)
multiples = array_ops.concat([
array_ops.ones([array_ops.rank(is_accepted)], dtype=dtypes.int32),
array_ops.shape(x)[independent_chain_ndims:],
], axis=0)
m = array_ops.tile(array_ops.reshape(is_accepted, expand_shape),
multiples)
m.set_shape(x.shape)
return m
def _check_mu(self, mu):
"""Return `mu` after validity checks and possibly with assertations."""
mu = ops.convert_to_tensor(mu)
cov = self._cov
if mu.dtype != cov.dtype:
raise TypeError(
"mu and cov must have the same dtype. Found mu.dtype = %s, " "cov.dtype = %s" % (mu.dtype, cov.dtype)
)
# Try to validate with static checks.
mu_shape = mu.get_shape()
cov_shape = cov.get_shape()
if mu_shape.is_fully_defined() and cov_shape.is_fully_defined():
if mu_shape != cov_shape[:-1]:
raise ValueError(
"mu.shape and cov.shape[:-1] should match. Found: mu.shape=%s, "
"cov.shape=%s" % (mu_shape, cov_shape)
)
else:
return mu
# Static checks could not be run, so possibly do dynamic checks.
if not self.validate_args:
return mu
else:
assert_same_rank = check_ops.assert_equal(
array_ops.rank(mu) + 1,
cov.rank(),
data=[
"mu should have rank 1 less than cov. Found: rank(mu) = ",
array_ops.rank(mu),
" rank(cov) = ",
cov.rank(),
],
)
with ops.control_dependencies([assert_same_rank]):
assert_same_shape = check_ops.assert_equal(
array_ops.shape(mu),
cov.vector_shape(),
data=[
"mu.shape and cov.shape[:-1] should match. " "Found: shape(mu) = ",
array_ops.shape(mu),
" shape(cov) = ",
cov.shape(),
],
)
return control_flow_ops.with_dependencies([assert_same_shape], mu)
def _remove_squeezable_dimensions(
labels, predictions, weights=None, expected_rank_diff=0):
"""Internal version of _remove_squeezable_dimensions which handles weights.
Squeezes `predictions` and `labels` if their ranks differ from expected by
exactly 1.
Squeezes `weights` if its rank is 1 more than the new rank of `predictions`
This will use static shape if available. Otherwise, it will add graph
operations, which could result in a performance hit.
Args:
labels: Label values, a `Tensor` whose dimensions match `predictions`.
predictions: Predicted values, a `Tensor` of arbitrary dimensions.
weights: Optional weight `Tensor`. It will be squeezed if it's not scalar,
and its rank is 1 more than the new rank of `labels`.
expected_rank_diff: Expected result of `rank(predictions) - rank(labels)`.
Returns:
Tuple of `predictions`, `labels` and `weights`, possibly with the last
dimension squeezed.
"""
labels, predictions = confusion_matrix.remove_squeezable_dimensions(
labels, predictions, expected_rank_diff=expected_rank_diff)
if weights is not None:
weights = ops.convert_to_tensor(weights)
labels_rank = labels.get_shape().ndims
weights_shape = weights.get_shape()
weights_rank = weights_shape.ndims
if (labels_rank is not None) and (weights_rank is not None):
# Use static rank.
rank_diff = weights_rank - labels_rank
if rank_diff == 1:
weights = array_ops.squeeze(weights, [-1])
return labels, predictions, weights
# Use dynamic rank.
rank_diff = array_ops.rank(weights) - array_ops.rank(labels)
if (weights_rank is None) or (
weights_rank > 0 and weights_shape.dims[-1].is_compatible_with(1)):
weights = control_flow_ops.cond(
math_ops.equal(1, rank_diff),
lambda: array_ops.squeeze(weights, [-1]),
lambda: weights)
return labels, predictions, weights
def sample(self, n, seed=None, name=None):
"""Sample `n` observations from the Multivariate Normal Distributions.
Args:
n: `Scalar`, type int32, the number of observations to sample.
seed: Python integer, the random seed.
name: The name to give this op.
Returns:
samples: `[n, ...]`, a `Tensor` of `n` samples for each
of the distributions determined by broadcasting the hyperparameters.
"""
with ops.op_scope(
[self._mu, self._sigma_chol, n], name, "MultivariateNormalSample"):
# TODO(ebrevdo): Is there a better way to get broadcast_shape?
broadcast_shape = self.mu.get_shape()
n = ops.convert_to_tensor(n)
sigma_shape_left = array_ops.slice(
array_ops.shape(self._sigma_chol),
[0], array_ops.pack([array_ops.rank(self._sigma_chol) - 2]))
k_n = array_ops.pack([self._k, n])
shape = array_ops.concat(0, [sigma_shape_left, k_n])
white_samples = random_ops.random_normal(
shape=shape, mean=0, stddev=1, dtype=self._mu.dtype, seed=seed)
correlated_samples = math_ops.batch_matmul(
self._sigma_chol, white_samples)
# Move the last dimension to the front
perm = array_ops.concat(
0,
(array_ops.pack([array_ops.rank(correlated_samples) - 1]),
math_ops.range(0, array_ops.rank(correlated_samples) - 1)))
# TODO(ebrevdo): Once we get a proper tensor contraction op,
# perform the inner product using that instead of batch_matmul
# and this slow transpose can go away!
correlated_samples = array_ops.transpose(correlated_samples, perm)
samples = correlated_samples + self.mu
# Provide some hints to shape inference
n_val = tensor_util.constant_value(n)
final_shape = tensor_shape.vector(n_val).concatenate(broadcast_shape)
samples.set_shape(final_shape)
return samples
def _hinge_loss(logits, target):
check_shape_op = control_flow_ops.Assert(
math_ops.less_equal(array_ops.rank(target), 2),
["target's shape should be either [batch_size, 1] or [batch_size]"])
with ops.control_dependencies([check_shape_op]):
target = array_ops.reshape(target, shape=[array_ops.shape(target)[0], 1])
return losses.hinge_loss(logits, target)
def _shuffle_to_front(input_tensor, k):
"""Shuffles the last `k` indices of `input_tensor` to the front.
Transposes `input_tensor` to have the last `k` indices at the front. The input
may have arbitrary rank and unknown shape.
Args:
input_tensor: A `Tensor` of arbitrary rank and unknown shape.
k: A scalar `Tensor` specifying how many indices to shuffle.
Returns:
A transposed version of `input_tensor` with `k` indices shuffled to the
front.
Raises:
ValueError: If `input_tensor` is not at least rank `k` or `k` is not scalar.
"""
k = ops.convert_to_tensor(k, name="k")
k.shape.with_rank(0)
k_static = tensor_util.constant_value(k)
if k_static is not None:
input_tensor.shape.with_rank_at_least(k_static)
rank = array_ops.rank(input_tensor)
outer_indices, inner_indices = array_ops.split(math_ops.range(rank),
[rank - k, k])
permutation = array_ops.concat([inner_indices, outer_indices], 0)
return array_ops.transpose(input_tensor, perm=permutation)
def _num_present(losses, weights, per_batch=False):
"""Computes the number of elements in the loss function induced by `weights`.
A given weights tensor induces different numbers of usable elements in the
`losses` tensor. The `weights` tensor is broadcast across `losses` for all
possible dimensions. For example, if `losses` is a tensor of dimension
`[4, 5, 6, 3]` and `weights` is a tensor of shape `[4, 5]`, then `weights` is,
in effect, tiled to match the shape of `losses`. Following this effective
tile, the total number of present elements is the number of non-zero weights.
Args:
losses: `Tensor` of shape `[batch_size, d1, ... dN]`.
weights: `Tensor` of shape `[]`, `[batch_size]` or
`[batch_size, d1, ... dK]`, where K < N.
per_batch: Whether to return the number of elements per batch or as a sum
total.
Returns:
The number of present (non-zero) elements in the losses tensor. If
`per_batch` is `True`, the value is returned as a tensor of size
`[batch_size]`. Otherwise, a single scalar tensor is returned.
"""
with ops.name_scope(None, "num_present", (losses, weights)) as scope:
weights = math_ops.to_float(weights)
present = array_ops.where(
math_ops.equal(weights, 0.0),
array_ops.zeros_like(weights),
array_ops.ones_like(weights))
present = weights_broadcast_ops.broadcast_weights(present, losses)
if per_batch:
return math_ops.reduce_sum(
present, axis=math_ops.range(1, array_ops.rank(present)),
keep_dims=True, name=scope)
return math_ops.reduce_sum(present, name=scope)
def rgb_to_grayscale(images, name=None):
"""Converts one or more images from RGB to Grayscale.
Outputs a tensor of the same `DType` and rank as `images`. The size of the
last dimension of the output is 1, containing the Grayscale value of the
pixels.
Args:
images: The RGB tensor to convert. Last dimension must have size 3 and
should contain RGB values.
name: A name for the operation (optional).
Returns:
The converted grayscale image(s).
"""
with ops.op_scope([images], name, 'rgb_to_grayscale') as name:
images = ops.convert_to_tensor(images, name='images')
# Remember original dtype to so we can convert back if needed
orig_dtype = images.dtype
flt_image = convert_image_dtype(images, dtypes.float32)
# Reference for converting between RGB and grayscale.
# https://en.wikipedia.org/wiki/Luma_%28video%29
rgb_weights = [0.2989, 0.5870, 0.1140]
rank_1 = array_ops.expand_dims(array_ops.rank(images) - 1, 0)
gray_float = math_ops.reduce_sum(flt_image * rgb_weights,
rank_1,
keep_dims=True)
gray_float.set_shape(images.get_shape()[:-1].concatenate([1]))
return convert_image_dtype(gray_float, orig_dtype, name=name)
def _softmax_cross_entropy_loss(logits, target):
check_shape_op = control_flow_ops.Assert(
math_ops.less_equal(array_ops.rank(target), 2),
["target's shape should be either [batch_size, 1] or [batch_size]"])
with ops.control_dependencies([check_shape_op]):
target = array_ops.reshape(target, shape=[array_ops.shape(target)[0]])
return nn.sparse_softmax_cross_entropy_with_logits(logits, target)
def _make_columnar(self, x):
"""Ensures non-scalar input has at least one column.
Example:
If `x = [1, 2, 3]` then the output is `[[1], [2], [3]]`.
If `x = [[1, 2, 3], [4, 5, 6]]` then the output is unchanged.
If `x = 1` then the output is unchanged.
Args:
x: `Tensor`.
Returns:
columnar_x: `Tensor` with at least two dimensions.
"""
if x.get_shape().ndims is not None:
if x.get_shape().ndims == 1:
x = x[array_ops.newaxis, :]
return x
shape = array_ops.shape(x)
maybe_expanded_shape = array_ops.concat([
shape[:-1],
distribution_util.pick_vector(
math_ops.equal(array_ops.rank(x), 1),
[1], np.array([], dtype=np.int32)),
shape[-1:],
], 0)
return array_ops.reshape(x, maybe_expanded_shape)
def _RightShift(x): """Shifts next-to-last dimension to the right, adding zero on the left.""" rank = array_ops.rank(x) zeros = array_ops.zeros((rank - 2, 2), dtype=dtypes.int32) pad = array_ops.concat([zeros, array_ops.constant([[1, 0], [0, 0]])], axis=0) return array_ops.pad(x[..., :-1, :], pad)
def __init__(
self,
temperature,
logits=None,
probs=None,
dtype=dtypes.float32,
validate_args=False,
allow_nan_stats=True,
name="ExpRelaxedOneHotCategorical"):
"""Initialize ExpRelaxedOneHotCategorical using class log-probabilities.
Args:
temperature: An 0-D `Tensor`, representing the temperature
of a set of ExpRelaxedCategorical distributions. The temperature should
be positive.
logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
of a set of ExpRelaxedCategorical distributions. The first
`N - 1` dimensions index into a batch of independent distributions and
the last dimension represents a vector of logits for each class. Only
one of `logits` or `probs` should be passed in.
probs: An N-D `Tensor`, `N >= 1`, representing the probabilities
of a set of ExpRelaxedCategorical distributions. The first
`N - 1` dimensions index into a batch of independent distributions and
the last dimension represents a vector of probabilities for each
class. Only one of `logits` or `probs` should be passed in.
dtype: The type of the event samples (default: int32).
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
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.
"""
parameters = locals()
with ops.name_scope(name, values=[logits, probs, temperature]):
with ops.control_dependencies([check_ops.assert_positive(temperature)]
if validate_args else []):
self._temperature = array_ops.identity(temperature, name="temperature")
self._temperature_2d = array_ops.reshape(temperature, [-1, 1],
name="temperature_2d")
self._logits, self._probs = distribution_util.get_logits_and_probs(
name=name, logits=logits, probs=probs, validate_args=validate_args,
multidimensional=True)
logits_shape_static = self._logits.get_shape().with_rank_at_least(1)
if logits_shape_static.ndims is not None:
self._batch_rank = ops.convert_to_tensor(
logits_shape_static.ndims - 1,
dtype=dtypes.int32,
name="batch_rank")
else:
with ops.name_scope(name="batch_rank"):
self._batch_rank = array_ops.rank(self._logits) - 1
with ops.name_scope(name="event_size"):
self._event_size = array_ops.shape(self._logits)[-1]
super(ExpRelaxedOneHotCategorical, self).__init__(
dtype=dtype,
reparameterization_type=distribution.FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
graph_parents=[self._logits,
self._probs,
self._temperature],
name=name)
def rotate_transpose(x, shift, name="rotate_transpose"):
"""Circularly moves dims left or right.
Effectively identical to:
```python
numpy.transpose(x, numpy.roll(numpy.arange(len(x.shape)), shift))
```
When `validate_args=False` additional graph-runtime checks are
performed. These checks entail moving data from to GPU to CPU.
Example:
```python
x = ... # Tensor of shape [1, 2, 3, 4].
rotate_transpose(x, -1) # result shape: [2, 3, 4, 1]
rotate_transpose(x, -2) # result shape: [3, 4, 1, 2]
rotate_transpose(x, 1) # result shape: [4, 1, 2, 3]
rotate_transpose(x, 2) # result shape: [3, 4, 1, 2]
rotate_transpose(x, 7) == rotate_transpose(x, 3)
rotate_transpose(x, -7) == rotate_transpose(x, -3)
```
Args:
x: `Tensor`.
shift: `Tensor`. Number of dimensions to transpose left (shift<0) or
transpose right (shift>0).
name: Python `str`. The name to give this op.
Returns:
rotated_x: Input `Tensor` with dimensions circularly rotated by shift.
Raises:
TypeError: if shift is not integer type.
"""
with ops.name_scope(name, values=[x, shift]):
x = ops.convert_to_tensor(x, name="x")
shift = ops.convert_to_tensor(shift, name="shift")
# We do not assign back to preserve constant-ness.
check_ops.assert_integer(shift)
shift_value_static = tensor_util.constant_value(shift)
ndims = x.get_shape().ndims
if ndims is not None and shift_value_static is not None:
if ndims < 2: return x
shift_value_static = np.sign(shift_value_static) * (
abs(shift_value_static) % ndims)
if shift_value_static == 0: return x
perm = np.roll(np.arange(ndims), shift_value_static)
return array_ops.transpose(x, perm=perm)
else:
# Consider if we always had a positive shift, and some specified
# direction.
# When shifting left we want the new array:
# last(x, n-shift) + first(x, shift)
# and if shifting right then we want:
# last(x, shift) + first(x, n-shift)
# Observe that last(a) == slice(a, n) and first(a) == slice(0, a).
# Also, we can encode direction and shift as one: direction * shift.
# Combining these facts, we have:
# a = cond(shift<0, -shift, n-shift)
# last(x, n-a) + first(x, a) == x[a:n] + x[0:a]
# Finally, we transform shift by modulo length so it can be specified
# independently from the array upon which it operates (like python).
ndims = array_ops.rank(x)
shift = array_ops.where(math_ops.less(shift, 0),
math_ops.mod(-shift, ndims),
ndims - math_ops.mod(shift, ndims))
first = math_ops.range(0, shift)
last = math_ops.range(shift, ndims)
perm = array_ops.concat([last, first], 0)
return array_ops.transpose(x, perm=perm)
def _mini_batch_training_op(self, inputs, cluster_idx_list, cluster_centers,
total_counts):
"""Creates an op for training for mini batch case.
Args:
inputs: list of input Tensors.
cluster_idx_list: A vector (or list of vectors). Each element in the
vector corresponds to an input row in 'inp' and specifies the cluster id
corresponding to the input.
cluster_centers: Tensor Ref of cluster centers.
total_counts: Tensor Ref of cluster counts.
Returns:
An op for doing an update of mini-batch k-means.
"""
update_ops = []
for inp, cluster_idx in zip(inputs, cluster_idx_list):
with ops.colocate_with(inp, ignore_existing=True):
assert total_counts is not None
cluster_idx = array_ops.reshape(cluster_idx, [-1])
# Dedupe the unique ids of cluster_centers being updated so that updates
# can be locally aggregated.
unique_ids, unique_idx = array_ops.unique(cluster_idx)
num_unique_cluster_idx = array_ops.size(unique_ids)
# Fetch the old values of counts and cluster_centers.
with ops.colocate_with(total_counts, ignore_existing=True):
old_counts = array_ops.gather(total_counts, unique_ids)
# TODO(agarwal): This colocation seems to run into problems. Fix it.
with ops.colocate_with(cluster_centers, ignore_existing=True):
old_cluster_centers = array_ops.gather(cluster_centers, unique_ids)
# Locally aggregate the increment to counts.
count_updates = math_ops.unsorted_segment_sum(
array_ops.ones_like(unique_idx, dtype=total_counts.dtype),
unique_idx, num_unique_cluster_idx)
# Locally compute the sum of inputs mapped to each id.
# For a cluster with old cluster value x, old count n, and with data
# d_1,...d_k newly assigned to it, we recompute the new value as
# \\(x += (sum_i(d_i) - k * x) / (n + k)\\).
# Compute \\(sum_i(d_i)\\), see comment above.
cluster_center_updates = math_ops.unsorted_segment_sum(
inp, unique_idx, num_unique_cluster_idx)
# Shape to enable broadcasting count_updates and learning_rate to inp.
# It extends the shape with 1's to match the rank of inp.
broadcast_shape = array_ops.concat([
array_ops.reshape(num_unique_cluster_idx, [1]),
array_ops.ones(
array_ops.reshape(array_ops.rank(inp) - 1, [1]),
dtype=dtypes.int32)
], 0)
# Subtract k * x, see comment above.
cluster_center_updates -= math_ops.cast(
array_ops.reshape(count_updates, broadcast_shape),
inp.dtype) * old_cluster_centers
learning_rate = math_ops.reciprocal(
math_ops.cast(old_counts + count_updates, inp.dtype))
learning_rate = array_ops.reshape(learning_rate, broadcast_shape)
# scale by 1 / (n + k), see comment above.
cluster_center_updates *= learning_rate
# Apply the updates.
update_counts = state_ops.scatter_add(total_counts, unique_ids,
count_updates)
update_cluster_centers = state_ops.scatter_add(
cluster_centers, unique_ids, cluster_center_updates)
update_ops.extend([update_counts, update_cluster_centers])
return control_flow_ops.group(*update_ops)
def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index):
"""Gradient for concat op.
Args:
op: An operation.
grad: `Tensor` or `IndexedSlices` representing the gradients with respect
to each output of the op.
start_value_index: An integer index of the first value in the op.inputs.
end_value_index: An integer index of the last value in the op.inputs.
dim_index: An interger index of concat_dim or axis parameter in op.inputs.
Returns:
Tensors representing the partial gradients with respect to each input
of the op.
Raises:
ValueError: if concat_dim/axis is not statically known.
"""
def _CreateDenseMaskAndBegin(sizes, concat_dim):
"""Create variables for iteratively slicing a dense gradients tensor."""
# Since shape is 1-D, shape_of_shape = [rank-of-inputs]
shape_of_shape = array_ops.shape(sizes[0])
# Make a vector of length equal to the input's dimensions,
# with 0's everywhere and 1 in the concat dim position.
# Note: Can't use sparse_to_dense since it isn't GPU-capable (for now)
mask = array_ops.concat([
array_ops.fill(array_ops.expand_dims(concat_dim, 0), 0), [1],
array_ops.fill(shape_of_shape - concat_dim - 1, 0)
], 0)
begin = array_ops.fill(shape_of_shape, 0)
return mask, begin
def _ExtractInputShapes(inputs):
"""Extract the shapes of a set of input tensors."""
if context.executing_eagerly():
return array_ops.shape_n(inputs)
sizes = []
fully_known = True
for x in inputs:
input_shape = array_ops.shape(x)
if not isinstance(input_shape,
ops.Tensor) or input_shape.op.type != "Const":
fully_known = False
break
sizes.append(input_shape)
if fully_known:
return sizes
else:
return array_ops.shape_n(inputs)
# Degenerate concatenation, just return grad.
if len(op.inputs) == 2:
return grad + [None] if end_value_index <= dim_index else [None] + grad
concat_dim = op.inputs[dim_index]
input_values = op.inputs[start_value_index:end_value_index]
out_grads = []
if isinstance(grad, ops.Tensor):
if context.executing_eagerly() or isinstance(concat_dim, ops.EagerTensor):
# Using mod here for convenience since concat_dim is already verified
# in concat implementation to be within the allowed [-rank, rank) range.
non_neg_concat_dim = (
concat_dim._numpy().item(0) % input_values[0]._rank()) # pylint: disable=protected-access
# All inputs are guaranteed to be EagerTensors in eager mode
sizes = pywrap_tensorflow.TFE_Py_TensorShapeSlice(input_values,
non_neg_concat_dim)
out_grads = array_ops.split(grad, sizes, non_neg_concat_dim)
else:
if constant_op.is_constant(concat_dim):
# If concat_dim is a constant defined in a different context,
# then we duplicate it in the current context to avoid passing it
# through an Enter node.
# This is a small optimization in general, but it is required when
# compiling with XLA, as XLA needs the concat input to be folded into a
# constant.
grad_context = control_flow_util.GetOutputContext(grad.op)
dim_context = control_flow_util.GetOutputContext(concat_dim.op)
if dim_context != grad_context:
value = tensor_util.constant_value(concat_dim)
concat_dim = constant_op.constant(value=value, dtype=concat_dim.dtype)
# Using mod here for convenience since concat_dim is already verified
# in concat implementation to be within the allowed [-rank, rank) range.
non_neg_concat_dim = concat_dim % array_ops.rank(input_values[0])
# Get the inputs' tensor shapes
sizes = _ExtractInputShapes(input_values)
# The magic number of 16 was found through benchmarking a range of sizes
# on CPUs and a Maxwell TitanX. A speedup was seen in a large majority of
# cases when switching implementations at N=16, but it is possible that
# there will be a small number of performance regressions.
if len(sizes) > 16:
# extract the size of each input along the concat dimension
sizes = array_ops.squeeze(
array_ops.slice(
array_ops.stack(sizes, axis=1), [non_neg_concat_dim, 0],
[1, -1]))
out_grads = array_ops.split(grad, sizes, non_neg_concat_dim)
else:
offset = gen_array_ops.concat_offset(non_neg_concat_dim, sizes)
for (begin, size) in zip(offset, sizes):
out_grads.append(array_ops.slice(grad, begin, size))
elif isinstance(grad, ops.IndexedSlices):
# Using mod here for convenience since concat_dim is already verified
# in concat implementation to be within the allowed [-rank, rank) range.
non_neg_concat_dim = concat_dim % array_ops.rank(input_values[0])
concat_dim_static = tensor_util.constant_value(concat_dim)
if concat_dim_static is None:
raise ValueError("Can only compute IndexedSlices gradient with "
"statically-known concat_dim")
if concat_dim_static < 0:
rank = tensor_util.constant_value(array_ops.rank(input_values[0]))
if rank is None:
raise ValueError("Can only compute IndexedSlices gradient with "
"negative concat_dim when first value rank is "
"statically-known.")
concat_dim_static %= rank
# Get the inputs' tensor shapes
sizes = [array_ops.shape(x) for x in input_values]
if concat_dim_static > 0:
# IndexedSlices, non_neg_concat_dim > 0. Each input gets IndexedSlices
# gradients with all the indices, but with grad.values sliced accordingly.
# This is like the Tensor case, except shape(grad.values)[0] is not equal
# to shape(sizes[i])[0], since only a subset of the dim-0 values are
# stored.
mask, begin = _CreateDenseMaskAndBegin(sizes, non_neg_concat_dim)
for size in sizes:
new_values = array_ops.slice(
grad.values, begin,
array_ops.concat([[-1], array_ops.slice(size, [1], [-1])], 0))
out_grads.append(ops.IndexedSlices(new_values, grad.indices, size))
# Lint complains begin = begin + ...
begin = math_ops.add(begin, size * mask)
else:
# IndexedSlices, concat_dim == 0. Each input gets IndexedSlices gradients
# only for the relevant indices.
start = constant_op.constant(0, dtype=grad.indices.dtype)
for size in sizes:
size_concat_dim = array_ops.gather(size, non_neg_concat_dim)
if size_concat_dim.dtype != grad.indices.dtype:
size_concat_dim = math_ops.cast(
size_concat_dim, dtype=grad.indices.dtype)
end = start + size_concat_dim
# Compute the 1-D Tensor of indices relevant for this input.
indices_to_select = array_ops.squeeze(
array_ops.where(
math_ops.logical_and(grad.indices >= start,
grad.indices < end)),
axis=[1])
new_indices = array_ops.gather(grad.indices, indices_to_select) - start
new_values = array_ops.gather(grad.values, indices_to_select)
out_grads.append(ops.IndexedSlices(new_values, new_indices, size))
start = end
else:
raise TypeError("Expected Tensor or IndexedSlices, got %s" % type(grad))
return (out_grads + [None]
if end_value_index <= dim_index else [None] + out_grads)
def update_confusion_matrix_variables(variables_to_update,
y_true,
y_pred,
thresholds,
top_k=None,
class_id=None,
sample_weight=None,
multi_label=False,
label_weights=None):
"""Returns op to update the given confusion matrix variables.
For every pair of values in y_true and y_pred:
true_positive: y_true == True and y_pred > thresholds
false_negatives: y_true == True and y_pred <= thresholds
true_negatives: y_true == False and y_pred <= thresholds
false_positive: y_true == False and y_pred > thresholds
The results will be weighted and added together. When multiple thresholds are
provided, we will repeat the same for every threshold.
For estimation of these metrics over a stream of data, the function creates an
`update_op` operation that updates the given variables.
If `sample_weight` is `None`, weights default to 1.
Use weights of 0 to mask values.
Args:
variables_to_update: Dictionary with 'tp', 'fn', 'tn', 'fp' as valid keys
and corresponding variables to update as values.
y_true: A `Tensor` whose shape matches `y_pred`. Will be cast to `bool`.
y_pred: A floating point `Tensor` of arbitrary shape and whose values are in
the range `[0, 1]`.
thresholds: A float value, float tensor, python list, or tuple of float
thresholds in `[0, 1]`, or NEG_INF (used when top_k is set).
top_k: Optional int, indicates that the positive labels should be limited to
the top k predictions.
class_id: Optional int, limits the prediction and labels to the class
specified by this argument.
sample_weight: Optional `Tensor` whose rank is either 0, or the same rank as
`y_true`, and must be broadcastable to `y_true` (i.e., all dimensions must
be either `1`, or the same as the corresponding `y_true` dimension).
multi_label: Optional boolean indicating whether multidimensional
prediction/labels should be treated as multilabel responses, or flattened
into a single label. When True, the valus of `variables_to_update` must
have a second dimension equal to the number of labels in y_true and
y_pred, and those tensors must not be RaggedTensors.
label_weights: (optional) tensor of non-negative weights for multilabel
data. The weights are applied when calculating TP, FP, FN, and TN without
explicit multilabel handling (i.e. when the data is to be flattened).
Returns:
Update op.
Raises:
ValueError: If `y_pred` and `y_true` have mismatched shapes, or if
`sample_weight` is not `None` and its shape doesn't match `y_pred`, or if
`variables_to_update` contains invalid keys.
"""
if multi_label and label_weights is not None:
raise ValueError(
'`label_weights` for multilabel data should be handled '
'outside of `update_confusion_matrix_variables` when '
'`multi_label` is True.')
if variables_to_update is None:
return
if not any(key
for key in variables_to_update if key in list(ConfusionMatrix)):
raise ValueError(
'Please provide at least one valid confusion matrix '
'variable to update. Valid variable key options are: "{}". '
'Received: "{}"'.format(list(ConfusionMatrix),
variables_to_update.keys()))
variable_dtype = list(variables_to_update.values())[0].dtype
y_true = math_ops.cast(y_true, dtype=variable_dtype)
y_pred = math_ops.cast(y_pred, dtype=variable_dtype)
thresholds = ops.convert_to_tensor_v2(thresholds, dtype=variable_dtype)
num_thresholds = thresholds.shape[0]
if multi_label:
one_thresh = math_ops.equal(math_ops.cast(1, dtype=dtypes.int32),
array_ops.rank(thresholds),
name='one_set_of_thresholds_cond')
else:
[y_pred, y_true
], _ = ragged_assert_compatible_and_get_flat_values([y_pred, y_true],
sample_weight)
one_thresh = math_ops.cast(True, dtype=dtypes.bool)
invalid_keys = [
key for key in variables_to_update if key not in list(ConfusionMatrix)
]
if invalid_keys:
raise ValueError(
'Invalid keys: {}. Valid variable key options are: "{}"'.format(
invalid_keys, list(ConfusionMatrix)))
with ops.control_dependencies([
check_ops.assert_greater_equal(y_pred,
math_ops.cast(0.0,
dtype=y_pred.dtype),
message='predictions must be >= 0'),
check_ops.assert_less_equal(y_pred,
math_ops.cast(1.0, dtype=y_pred.dtype),
message='predictions must be <= 1')
]):
if sample_weight is None:
y_pred, y_true = tf_losses_utils.squeeze_or_expand_dimensions(
y_pred, y_true)
else:
y_pred, y_true, sample_weight = (
tf_losses_utils.squeeze_or_expand_dimensions(
y_pred, y_true, sample_weight=sample_weight))
y_pred.shape.assert_is_compatible_with(y_true.shape)
if top_k is not None:
y_pred = _filter_top_k(y_pred, top_k)
if class_id is not None:
y_true = y_true[..., class_id]
y_pred = y_pred[..., class_id]
pred_shape = array_ops.shape(y_pred)
num_predictions = pred_shape[0]
if y_pred.shape.ndims == 1:
num_labels = 1
else:
num_labels = gen_math_ops.Prod(input=pred_shape[1:], axis=0)
thresh_label_tile = control_flow_ops.cond(
one_thresh, lambda: num_labels,
lambda: math_ops.cast(1, dtype=dtypes.int32))
# Reshape predictions and labels, adding a dim for thresholding.
if multi_label:
predictions_extra_dim = array_ops.expand_dims(y_pred, 0)
labels_extra_dim = array_ops.expand_dims(
math_ops.cast(y_true, dtype=dtypes.bool), 0)
else:
# Flatten predictions and labels when not multilabel.
predictions_extra_dim = array_ops.reshape(y_pred, [1, -1])
labels_extra_dim = array_ops.reshape(
math_ops.cast(y_true, dtype=dtypes.bool), [1, -1])
# Tile the thresholds for every prediction.
if multi_label:
thresh_pretile_shape = [num_thresholds, 1, -1]
thresh_tiles = [1, num_predictions, thresh_label_tile]
data_tiles = [num_thresholds, 1, 1]
else:
thresh_pretile_shape = [num_thresholds, -1]
thresh_tiles = [1, num_predictions * num_labels]
data_tiles = [num_thresholds, 1]
thresh_tiled = array_ops.tile(
array_ops.reshape(thresholds, thresh_pretile_shape),
array_ops.stack(thresh_tiles))
# Tile the predictions for every threshold.
preds_tiled = array_ops.tile(predictions_extra_dim, data_tiles)
# Compare predictions and threshold.
pred_is_pos = math_ops.greater(preds_tiled, thresh_tiled)
# Tile labels by number of thresholds
label_is_pos = array_ops.tile(labels_extra_dim, data_tiles)
if sample_weight is not None:
sample_weight = weights_broadcast_ops.broadcast_weights(
math_ops.cast(sample_weight, dtype=variable_dtype), y_pred)
weights_tiled = array_ops.tile(
array_ops.reshape(sample_weight, thresh_tiles), data_tiles)
else:
weights_tiled = None
if label_weights is not None and not multi_label:
label_weights = array_ops.expand_dims(label_weights, 0)
label_weights = weights_broadcast_ops.broadcast_weights(
label_weights, y_pred)
label_weights_tiled = array_ops.tile(
array_ops.reshape(label_weights, thresh_tiles), data_tiles)
if weights_tiled is None:
weights_tiled = label_weights_tiled
else:
weights_tiled = math_ops.multiply(weights_tiled,
label_weights_tiled)
update_ops = []
def weighted_assign_add(label, pred, weights, var):
label_and_pred = math_ops.cast(math_ops.logical_and(label, pred),
dtype=var.dtype)
if weights is not None:
label_and_pred *= math_ops.cast(weights, dtype=var.dtype)
return var.assign_add(math_ops.reduce_sum(label_and_pred, 1))
loop_vars = {
ConfusionMatrix.TRUE_POSITIVES: (label_is_pos, pred_is_pos),
}
update_tn = ConfusionMatrix.TRUE_NEGATIVES in variables_to_update
update_fp = ConfusionMatrix.FALSE_POSITIVES in variables_to_update
update_fn = ConfusionMatrix.FALSE_NEGATIVES in variables_to_update
if update_fn or update_tn:
pred_is_neg = math_ops.logical_not(pred_is_pos)
loop_vars[ConfusionMatrix.FALSE_NEGATIVES] = (label_is_pos,
pred_is_neg)
if update_fp or update_tn:
label_is_neg = math_ops.logical_not(label_is_pos)
loop_vars[ConfusionMatrix.FALSE_POSITIVES] = (label_is_neg,
pred_is_pos)
if update_tn:
loop_vars[ConfusionMatrix.TRUE_NEGATIVES] = (label_is_neg,
pred_is_neg)
for matrix_cond, (label, pred) in loop_vars.items():
if matrix_cond in variables_to_update:
update_ops.append(
weighted_assign_add(label, pred, weights_tiled,
variables_to_update[matrix_cond]))
return control_flow_ops.group(update_ops)
def _grad(op, grad):
"""A gradient function for RFFT with the provided `rank` and `irfft_fn`."""
fft_length = op.inputs[1]
complex_dtype = grad.dtype
real_dtype = complex_dtype.real_dtype
input_shape = _array_ops.shape(op.inputs[0])
is_even = _math_ops.cast(1 - (fft_length[-1] % 2), complex_dtype)
def _tile_for_broadcasting(matrix, t):
expanded = _array_ops.reshape(
matrix,
_array_ops.concat([
_array_ops.ones([_array_ops.rank(t) - 2], _dtypes.int32),
_array_ops.shape(matrix)
], 0))
return _array_ops.tile(
expanded,
_array_ops.concat([_array_ops.shape(t)[:-2], [1, 1]], 0))
def _mask_matrix(length):
"""Computes t_n = exp(sqrt(-1) * pi * n^2 / line_len)."""
# TODO(rjryan): Speed up computation of twiddle factors using the
# following recurrence relation and cache them across invocations of RFFT.
#
# t_n = exp(sqrt(-1) * pi * n^2 / line_len)
# for n = 0, 1,..., line_len-1.
# For n > 2, use t_n = t_{n-1}^2 / t_{n-2} * t_1^2
a = _array_ops.tile(
_array_ops.expand_dims(_math_ops.range(length), 0),
(length, 1))
b = _array_ops.transpose(a, [1, 0])
return _math_ops.exp(-2j * np.pi *
_math_ops.cast(a * b, complex_dtype) /
_math_ops.cast(length, complex_dtype))
def _ymask(length):
"""A sequence of [1+0j, -1+0j, 1+0j, -1+0j, ...] with length `length`."""
return _math_ops.cast(1 - 2 * (_math_ops.range(length) % 2),
complex_dtype)
y0 = grad[..., 0:1]
if rank == 1:
ym = grad[..., -1:]
extra_terms = y0 + is_even * ym * _ymask(input_shape[-1])
elif rank == 2:
# Create a mask matrix for y0 and ym.
base_mask = _mask_matrix(input_shape[-2])
# Tile base_mask to match y0 in shape so that we can batch-matmul the
# inner 2 dimensions.
tiled_mask = _tile_for_broadcasting(base_mask, y0)
y0_term = _math_ops.matmul(tiled_mask, _math_ops.conj(y0))
extra_terms = y0_term
ym = grad[..., -1:]
ym_term = _math_ops.matmul(tiled_mask, _math_ops.conj(ym))
inner_dim = input_shape[-1]
ym_term = _array_ops.tile(
ym_term,
_array_ops.concat([
_array_ops.ones([_array_ops.rank(grad) - 1],
_dtypes.int32), [inner_dim]
], 0)) * _ymask(inner_dim)
extra_terms += is_even * ym_term
# The gradient of RFFT is the IRFFT of the incoming gradient times a scaling
# factor, plus some additional terms to make up for the components dropped
# due to Hermitian symmetry.
input_size = _math_ops.cast(_fft_size_for_grad(op.inputs[0], rank),
real_dtype)
the_irfft = irfft_fn(grad, fft_length)
return 0.5 * (the_irfft * input_size +
_math_ops.real(extra_terms)), None
def __init__(self,
logits=None,
p=None,
dtype=dtypes.int32,
validate_args=False,
allow_nan_stats=True,
name="Categorical"):
"""Initialize Categorical distributions using class log-probabilities.
Args:
logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
of a set of Categorical distributions. The first `N - 1` dimensions
index into a batch of independent distributions and the last dimension
represents a vector of logits for each class. Only one of `logits` or
`p` should be passed in.
p: An N-D `Tensor`, `N >= 1`, representing the probabilities
of a set of Categorical distributions. The first `N - 1` dimensions
index into a batch of independent distributions and the last dimension
represents a vector of probabilities for each class. Only one of
`logits` or `p` should be passed in.
dtype: The type of the event samples (default: int32).
validate_args: Unused in this distribution.
allow_nan_stats: `Boolean`, default `True`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
name: A name for this distribution (optional).
"""
parameters = locals()
parameters.pop("self")
with ops.name_scope(name, values=[logits]) as ns:
self._logits, self._p = distribution_util.get_logits_and_prob(
name=name,
logits=logits,
p=p,
validate_args=validate_args,
multidimensional=True)
logits_shape_static = self._logits.get_shape().with_rank_at_least(
1)
if logits_shape_static.ndims is not None:
self._batch_rank = ops.convert_to_tensor(
logits_shape_static.ndims - 1,
dtype=dtypes.int32,
name="batch_rank")
else:
with ops.name_scope(name="batch_rank"):
self._batch_rank = array_ops.rank(self._logits) - 1
logits_shape = array_ops.shape(self._logits, name="logits_shape")
if logits_shape_static[-1].value is not None:
self._num_classes = ops.convert_to_tensor(
logits_shape_static[-1].value,
dtype=dtypes.int32,
name="num_classes")
else:
self._num_classes = array_ops.gather(logits_shape,
self._batch_rank,
name="num_classes")
if logits_shape_static[:-1].is_fully_defined():
self._batch_shape_val = constant_op.constant(
logits_shape_static[:-1].as_list(),
dtype=dtypes.int32,
name="batch_shape")
else:
with ops.name_scope(name="batch_shape"):
self._batch_shape_val = logits_shape[:-1]
super(Categorical,
self).__init__(dtype=dtype,
is_continuous=False,
is_reparameterized=False,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
graph_parents=[self._logits, self._num_classes],
name=ns)
def embedding_lookup_sparse(params,
sp_ids,
sp_weights,
partition_strategy="mod",
name=None,
combiner=None,
max_norm=None):
"""Looks up embeddings for the given ids and weights from a list of tensors.
This op assumes that there is at least one id for each row in the dense tensor
represented by sp_ids (i.e. there are no rows with empty features), and that
all the indices of sp_ids are in canonical row-major order.
`sp_ids` and `sp_weights` (if not None) are `SparseTensor`s with rank of 2.
Embeddings are always aggregated along the last dimension.
It also assumes that all id values lie in the range [0, p0), where p0
is the sum of the size of params along dimension 0.
Args:
params: A single tensor representing the complete embedding tensor, or a
list tensors all of same shape except for the first dimension,
representing sharded embedding tensors. Alternatively, a
`PartitionedVariable`, created by partitioning along dimension 0. Each
element must be appropriately sized for the given `partition_strategy`.
sp_ids: N x M `SparseTensor` of int64 ids where N is typically batch size
and M is arbitrary.
sp_weights: either a `SparseTensor` of float / double weights, or `None` to
indicate all weights should be taken to be 1. If specified, `sp_weights`
must have exactly the same shape and indices as `sp_ids`.
partition_strategy: A string specifying the partitioning strategy, relevant
if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
is `"mod"`. See `tf.nn.embedding_lookup` for more details.
name: Optional name for the op.
combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
and "sum" are supported. "sum" computes the weighted sum of the embedding
results for each row. "mean" is the weighted sum divided by the total
weight. "sqrtn" is the weighted sum divided by the square root of the sum
of the squares of the weights. Defaults to `mean`.
max_norm: If not `None`, each embedding is clipped if its l2-norm is larger
than this value, before combining.
Returns:
A dense tensor representing the combined embeddings for the
sparse ids. For each row in the dense tensor represented by `sp_ids`, the op
looks up the embeddings for all ids in that row, multiplies them by the
corresponding weight, and combines these embeddings as specified.
In other words, if
`shape(combined params) = [p0, p1, ..., pm]`
and
`shape(sp_ids) = shape(sp_weights) = [d0, d1]`
then
`shape(output) = [d0, p1, ..., pm]`.
For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are
```python
[0, 0]: id 1, weight 2.0
[0, 1]: id 3, weight 0.5
[1, 0]: id 0, weight 1.0
[2, 3]: id 1, weight 3.0
```
with `combiner`="mean", then the output will be a 3x20 matrix where
```python
output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
output[1, :] = (params[0, :] * 1.0) / 1.0
output[2, :] = (params[1, :] * 3.0) / 3.0
```
Raises:
TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is
neither `None` nor `SparseTensor`.
ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}.
"""
if combiner is None:
combiner = "mean"
if combiner not in ("mean", "sqrtn", "sum"):
raise ValueError(
f"combiner must be one of 'mean', 'sqrtn' or 'sum', got {combiner}")
if isinstance(params, variables.PartitionedVariable):
params = list(params) # Iterate to get the underlying Variables.
if not isinstance(params, list):
params = [params]
if not isinstance(sp_ids, sparse_tensor.SparseTensor):
raise TypeError(f"sp_ids must be SparseTensor, got {type(sp_ids)}")
ignore_weights = sp_weights is None
if not ignore_weights:
if not isinstance(sp_weights, sparse_tensor.SparseTensor):
raise TypeError(f"sp_weights must be either None or SparseTensor,"
f"got {type(sp_weights)}")
sp_ids.values.get_shape().assert_is_compatible_with(
sp_weights.values.get_shape())
sp_ids.indices.get_shape().assert_is_compatible_with(
sp_weights.indices.get_shape())
sp_ids.dense_shape.get_shape().assert_is_compatible_with(
sp_weights.dense_shape.get_shape())
# TODO(yleon): Add enhanced node assertions to verify that sp_ids and
# sp_weights have equal indices and shapes.
with ops.name_scope(name, "embedding_lookup_sparse",
params + [sp_ids]) as name:
segment_ids = sp_ids.indices[:, 0]
ids = sp_ids.values
ids, idx = array_ops.unique(ids)
embeddings = embedding_lookup(
params, ids, partition_strategy=partition_strategy, max_norm=max_norm)
if not ignore_weights:
if segment_ids.dtype != dtypes.int32:
segment_ids = math_ops.cast(segment_ids, dtypes.int32)
weights = sp_weights.values
embeddings = array_ops.gather(embeddings, idx)
original_dtype = embeddings.dtype
if embeddings.dtype in (dtypes.float16, dtypes.bfloat16):
# Cast low-precision embeddings to float32 during the computation to
# avoid numerical issues.
embeddings = math_ops.cast(embeddings, dtypes.float32)
if weights.dtype != embeddings.dtype:
weights = math_ops.cast(weights, embeddings.dtype)
# Reshape weights to allow broadcast
ones_shape = array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0)
ones = array_ops.ones(ones_shape, dtype=dtypes.int32)
bcast_weights_shape = array_ops.concat([array_ops.shape(weights), ones],
0)
orig_weights_shape = weights.get_shape()
weights = array_ops.reshape(weights, bcast_weights_shape)
# Set the weight shape, since after reshaping to bcast_weights_shape,
# the shape becomes None.
if embeddings.get_shape().ndims is not None:
weights.set_shape(
orig_weights_shape.concatenate(
[1 for _ in range(embeddings.get_shape().ndims - 1)]))
embeddings *= weights
if combiner == "sum":
embeddings = math_ops.segment_sum(embeddings, segment_ids, name=name)
elif combiner == "mean":
embeddings = math_ops.segment_sum(embeddings, segment_ids)
weight_sum = math_ops.segment_sum(weights, segment_ids)
embeddings = math_ops.div_no_nan(embeddings, weight_sum, name=name)
elif combiner == "sqrtn":
embeddings = math_ops.segment_sum(embeddings, segment_ids)
weights_squared = math_ops.pow(weights, 2)
weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
weight_sum_sqrt = math_ops.sqrt(weight_sum)
embeddings = math_ops.div_no_nan(embeddings, weight_sum_sqrt, name=name)
else:
assert False, "Unrecognized combiner"
if embeddings.dtype != original_dtype:
embeddings = math_ops.cast(embeddings, original_dtype)
else:
if segment_ids.dtype not in (dtypes.int32, dtypes.int64):
segment_ids = math_ops.cast(segment_ids, dtypes.int32)
assert idx is not None
if combiner == "sum":
embeddings = math_ops.sparse_segment_sum(
embeddings, idx, segment_ids, name=name)
elif combiner == "mean":
embeddings = math_ops.sparse_segment_mean(
embeddings, idx, segment_ids, name=name)
elif combiner == "sqrtn":
embeddings = math_ops.sparse_segment_sqrt_n(
embeddings, idx, segment_ids, name=name)
else:
assert False, "Unrecognized combiner"
return embeddings
def sign_magnitude_positive_definite(raw,
off_diagonal_scale=0.,
overall_scale=0.):
"""Constructs a positive definite matrix from an unconstrained input matrix.
We want to keep the whole matrix on a log scale, but also allow off-diagonal
elements to be negative, so the sign of off-diagonal elements is modeled
separately from their magnitude (using the lower and upper triangles
respectively). Specifically:
for i < j, we have:
output_cholesky[i, j] = raw[j, i] / (abs(raw[j, i]) + 1) *
exp((off_diagonal_scale + overall_scale + raw[i, j]) / 2)
output_cholesky[i, i] = exp((raw[i, i] + overall_scale) / 2)
output = output_cholesky^T * output_cholesky
where raw, off_diagonal_scale, and overall_scale are
un-constrained real-valued variables. The resulting values are stable
around zero due to the exponential (and the softsign keeps the function
smooth).
Args:
raw: A [..., M, M] Tensor.
off_diagonal_scale: A scalar or [...] shaped Tensor controlling the relative
scale of off-diagonal values in the output matrix.
overall_scale: A scalar or [...] shaped Tensor controlling the overall scale
of the output matrix.
Returns:
The `output` matrix described above, a [..., M, M] positive definite matrix.
"""
raw = ops.convert_to_tensor(raw)
diagonal = array_ops.matrix_diag_part(raw)
def _right_pad_with_ones(tensor, target_rank):
# Allow broadcasting even if overall_scale and off_diagonal_scale have batch
# dimensions
tensor = ops.convert_to_tensor(tensor, dtype=raw.dtype.base_dtype)
return array_ops.reshape(
tensor,
array_ops.concat([
array_ops.shape(tensor),
array_ops.ones([target_rank - array_ops.rank(tensor)],
dtype=target_rank.dtype)
],
axis=0))
# We divide the log values by 2 to compensate for the squaring that happens
# when transforming Cholesky factors into positive definite matrices.
sign_magnitude = (gen_math_ops.exp(
(raw + _right_pad_with_ones(off_diagonal_scale, array_ops.rank(raw)) +
_right_pad_with_ones(overall_scale, array_ops.rank(raw))) / 2.) *
nn.softsign(array_ops.matrix_transpose(raw)))
sign_magnitude.set_shape(raw.get_shape())
cholesky_factor = array_ops.matrix_set_diag(
input=array_ops.matrix_band_part(sign_magnitude, 0, -1),
diagonal=gen_math_ops.exp(
(diagonal +
_right_pad_with_ones(overall_scale, array_ops.rank(diagonal))) /
2.))
return math_ops.matmul(cholesky_factor, cholesky_factor, transpose_a=True)
def _padding_for_dimension(data, axis, pad_value):
"""Tile `pad_value` so it can be used to pad `data` at the given axis.
Returns a tensor `result` that has the same shape as `data` up to dimension
`axis`, but where each value `data[i0...iaxis]` is replaced by `pad_value`.
I.e., returns `result[i0...iaxis, j0...jN] = pad_value[j0...jN]`
(where `N=rank(pad_value)`).
Args:
data: The potentially ragged tensor that will be padded.
axis: The axis along which padding will be added.
pad_value: The padding value that should be used, or None if no padding will
be added. `rank(pad_value)` must be `rank(data) - axis`, and
`pad_value.shape[1:]` must be compatible with `data.shape[axis + 1:]`.
Returns:
A padding tensor with the same rank as `data`, which can be concatenated
to `data` to add padding.
"""
if pad_value is None:
return None
# Verify shape compatibility.
pad_value.shape[1:].assert_is_compatible_with(data.shape[axis:][1:])
if not isinstance(data, ragged_tensor.RaggedTensor):
data_shape = array_ops.shape(data)
pad_shape = array_ops.shape(pad_value)
outer_dimensions = data_shape[:axis]
expanded_pad_shape = array_ops.concat(
[array_ops.ones_like(outer_dimensions), pad_shape], axis=0)
tile_multiples = array_ops.concat(
[outer_dimensions,
array_ops.ones_like(pad_shape)], axis=0)
tiled_padding = array_ops.tile(
array_ops.reshape(pad_value, expanded_pad_shape), tile_multiples)
tiled_padding.set_shape(data.shape[:axis].concatenate(pad_value.shape))
return tiled_padding
assert axis >= 0
# Return the padding as-is if we're padding the outermost dimension.
if axis == 0:
return pad_value
elif axis == 1:
if isinstance(pad_value, ragged_tensor.RaggedTensor):
pad_rank = array_ops.rank(
pad_value.flat_values) + pad_value.ragged_rank
pad_nrows = pad_value.nrows()
else:
pad_rank = array_ops.rank(pad_value)
pad_nrows = array_ops.shape(pad_value, out_type=dtypes.int64)[0]
# Return a RaggedTensor that has the same number of rows as `data`, where
# each row contains a single copy of `pad_value`.
data_nrows = data.nrows()
pad_repeats = array_ops.concat(
[[math_ops.cast(data_nrows, dtypes.int32)],
array_ops.ones([pad_rank - 1], dtypes.int32)],
axis=0)
result_values = array_ops.tile(pad_value, pad_repeats)
return ragged_tensor.RaggedTensor.from_row_splits(
result_values,
math_ops.range(0, data_nrows + 1) * pad_nrows)
else: # Recurse if axis>1.
return ragged_tensor.RaggedTensor.from_row_splits(
_padding_for_dimension(data.values, axis - 1, pad_value),
data.row_splits)
def proposal_log_prob(x):
event_dims = math_ops.range(independent_chain_ndims,
array_ops.rank(x))
return -0.5 * math_ops.reduce_sum(x**2. + np.log(2 * np.pi),
axis=event_dims)
def _log_prob(self, x):
if self.cholesky_input_output_matrices:
x_sqrt = x
else:
# Complexity: O(nbk^3)
x_sqrt = linalg_ops.cholesky(x)
batch_shape = self.batch_shape_tensor()
event_shape = self.event_shape_tensor()
ndims = array_ops.rank(x_sqrt)
# sample_ndims = ndims - batch_ndims - event_ndims
sample_ndims = ndims - array_ops.shape(batch_shape)[0] - 2
sample_shape = array_ops.strided_slice(
array_ops.shape(x_sqrt), [0], [sample_ndims])
# We need to be able to pre-multiply each matrix by its corresponding
# batch scale matrix. Since a Distribution Tensor supports multiple
# samples per batch, this means we need to reshape the input matrix `x`
# so that the first b dimensions are batch dimensions and the last two
# are of shape [dimension, dimensions*number_of_samples]. Doing these
# gymnastics allows us to do a batch_solve.
#
# After we're done with sqrt_solve (the batch operation) we need to undo
# this reshaping so what we're left with is a Tensor partitionable by
# sample, batch, event dimensions.
# Complexity: O(nbk**2) since transpose must access every element.
scale_sqrt_inv_x_sqrt = x_sqrt
perm = array_ops.concat([math_ops.range(sample_ndims, ndims),
math_ops.range(0, sample_ndims)], 0)
scale_sqrt_inv_x_sqrt = array_ops.transpose(scale_sqrt_inv_x_sqrt, perm)
shape = array_ops.concat(
(batch_shape, (math_ops.cast(
self.dimension, dtype=dtypes.int32), -1)),
0)
scale_sqrt_inv_x_sqrt = array_ops.reshape(scale_sqrt_inv_x_sqrt, shape)
# Complexity: O(nbM*k) where M is the complexity of the operator solving
# a vector system. E.g., for OperatorPDDiag, each solve is O(k), so
# this complexity is O(nbk**2). For OperatorPDCholesky, each solve is
# O(k**2) so this step has complexity O(nbk^3).
scale_sqrt_inv_x_sqrt = self.scale_operator_pd.sqrt_solve(
scale_sqrt_inv_x_sqrt)
# Undo make batch-op ready.
# Complexity: O(nbk**2)
shape = array_ops.concat([batch_shape, event_shape, sample_shape], 0)
scale_sqrt_inv_x_sqrt = array_ops.reshape(scale_sqrt_inv_x_sqrt, shape)
perm = array_ops.concat([math_ops.range(ndims - sample_ndims, ndims),
math_ops.range(0, ndims - sample_ndims)], 0)
scale_sqrt_inv_x_sqrt = array_ops.transpose(scale_sqrt_inv_x_sqrt, perm)
# Write V = SS', X = LL'. Then:
# tr[inv(V) X] = tr[inv(S)' inv(S) L L']
# = tr[inv(S) L L' inv(S)']
# = tr[(inv(S) L) (inv(S) L)']
# = sum_{ik} (inv(S) L)_{ik}**2
# The second equality follows from the cyclic permutation property.
# Complexity: O(nbk**2)
trace_scale_inv_x = math_ops.reduce_sum(
math_ops.square(scale_sqrt_inv_x_sqrt),
axis=[-2, -1])
# Complexity: O(nbk)
half_log_det_x = math_ops.reduce_sum(
math_ops.log(array_ops.matrix_diag_part(x_sqrt)),
axis=[-1])
# Complexity: O(nbk**2)
log_prob = ((self.df - self.dimension - 1.) * half_log_det_x -
0.5 * trace_scale_inv_x -
self.log_normalization())
# Set shape hints.
# Try to merge what we know from the input then what we know from the
# parameters of this distribution.
if x.get_shape().ndims is not None:
log_prob.set_shape(x.get_shape()[:-2])
if (log_prob.get_shape().ndims is not None and
self.batch_shape.ndims is not None and
self.batch_shape.ndims > 0):
log_prob.get_shape()[-self.batch_shape.ndims:].merge_with(
self.batch_shape)
return log_prob
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops.rank(x_i)
def _event_shape(self):
return array_ops.gather(array_ops.shape(self._mean_val),
[array_ops.rank(self._mean_val) - 1])
def _process_input_helper(self,
update_row_factors,
sp_input=None,
transpose_input=False,
row_weights=None):
"""Creates the graph for processing a sparse slice of input.
Args:
update_row_factors: if True, update or project the row_factors, else
update or project the column factors.
sp_input: Please refer to comments for update_row_factors,
update_col_factors, project_row_factors, and project_col_factors for
restrictions.
transpose_input: If True, the input is logically transposed and then the
corresponding rows/columns of the transposed input are updated.
row_weights: If not None, this is the row/column weights to be used for
the update or projection. If None, use the corresponding weights from
the model. Note that the feature (column/row) weights will be
determined by the model. When not None, it can either be a scalar or
a rank-1 tensor with the same number of elements as the number of rows
of columns to be updated/projected.
Returns:
A tuple consisting of the following elements:
new_values: New values for the row/column factors.
update_op: An op that assigns the newly computed values to the row/column
factors.
unregularized_loss: A tensor (scalar) that contains the normalized
minibatch loss corresponding to sp_input, without the regularization
term. Add the regularization term below to yield the loss.
regularization: A tensor (scalar) that contains the normalized
regularization term for the minibatch loss corresponding to sp_input.
sum_weights: The sum of the weights corresponding to sp_input. This
can be used with unregularized loss to calculate the root weighted
squared error.
"""
assert isinstance(sp_input, sparse_tensor.SparseTensor)
if update_row_factors:
left = self._row_factors
right_factors = self._col_factors_cache
row_wt = self._row_wt_cache
col_wt = self._col_wt_cache
total_rows = self._input_rows
total_cols = self._input_cols
sharding_func = WALSModel._get_sharding_func(
self._input_rows, self._num_row_shards)
gramian = self._col_gramian_cache
else:
left = self._col_factors
right_factors = self._row_factors_cache
row_wt = self._col_wt_cache
col_wt = self._row_wt_cache
total_rows = self._input_cols
total_cols = self._input_rows
sharding_func = WALSModel._get_sharding_func(
self._input_cols, self._num_col_shards)
gramian = self._row_gramian_cache
transpose_input = not transpose_input
# Note that the row indices of sp_input are based on the original full input
# Here we reindex the rows and give them contiguous ids starting at 0.
# We use tf.unique to achieve this reindexing. Note that this is done so
# that the downstream kernel can assume that the input is "dense" along the
# row dimension.
row_ids, col_ids = array_ops.split(value=sp_input.indices,
num_or_size_splits=2,
axis=1)
update_row_indices, all_row_ids = array_ops.unique(row_ids[:, 0])
update_col_indices, all_col_ids = array_ops.unique(col_ids[:, 0])
col_ids = array_ops.expand_dims(
math_ops.cast(all_col_ids, dtypes.int64), 1)
row_ids = array_ops.expand_dims(
math_ops.cast(all_row_ids, dtypes.int64), 1)
if transpose_input:
update_indices = update_col_indices
row_shape = [
math_ops.cast(
array_ops.shape(update_row_indices)[0], dtypes.int64)
]
gather_indices = update_row_indices
else:
update_indices = update_row_indices
row_shape = [
math_ops.cast(
array_ops.shape(update_col_indices)[0], dtypes.int64)
]
gather_indices = update_col_indices
num_rows = math_ops.cast(
array_ops.shape(update_indices)[0], dtypes.int64)
col_shape = [num_rows]
right = embedding_ops.embedding_lookup(right_factors,
gather_indices,
partition_strategy="div")
new_sp_indices = array_ops.concat([row_ids, col_ids], 1)
new_sp_shape = (array_ops.concat([row_shape, col_shape], 0)
if transpose_input else array_ops.concat(
[col_shape, row_shape], 0))
new_sp_input = sparse_tensor.SparseTensor(indices=new_sp_indices,
values=sp_input.values,
dense_shape=new_sp_shape)
# Compute lhs and rhs of the normal equations
total_lhs = (self._unobserved_weight * gramian)
if self._regularization_matrix is not None:
total_lhs += self._regularization_matrix
if self._row_weights is None:
# Special case of ALS. Use a much simpler update rule.
total_rhs = (self._unobserved_weight *
sparse_ops.sparse_tensor_dense_matmul(
new_sp_input, right, adjoint_a=transpose_input))
# TODO(rmlarsen): handle transposing in tf.matrix_solve instead of
# transposing explicitly.
# TODO(rmlarsen): multi-thread tf.matrix_solve.
new_left_values = array_ops.transpose(
linalg_ops.matrix_solve(total_lhs,
array_ops.transpose(total_rhs)))
else:
if row_weights is None:
# TODO(yifanchen): Add special handling for single shard without using
# embedding_lookup and perform benchmarks for those cases. Same for
# col_weights lookup below.
row_weights_slice = embedding_ops.embedding_lookup(
row_wt, update_indices, partition_strategy="div")
else:
num_indices = array_ops.shape(update_indices)[0]
with ops.control_dependencies([
check_ops.assert_less_equal(
array_ops.rank(row_weights), 1)
]):
row_weights_slice = control_flow_ops.cond(
math_ops.equal(array_ops.rank(row_weights), 0), lambda:
(array_ops.ones([num_indices]) * row_weights),
lambda: math_ops.cast(row_weights, dtypes.float32))
col_weights = embedding_ops.embedding_lookup(
col_wt, gather_indices, partition_strategy="div")
partial_lhs, total_rhs = (
gen_factorization_ops.wals_compute_partial_lhs_and_rhs(
right,
col_weights,
self._unobserved_weight,
row_weights_slice,
new_sp_input.indices,
new_sp_input.values,
num_rows,
transpose_input,
name="wals_compute_partial_lhs_rhs"))
total_lhs = array_ops.expand_dims(total_lhs, 0) + partial_lhs
total_rhs = array_ops.expand_dims(total_rhs, -1)
new_left_values = array_ops.squeeze(
linalg_ops.matrix_solve(total_lhs, total_rhs), [2])
update_op_name = "row_update" if update_row_factors else "col_update"
update_op = self.scatter_update(left,
update_indices,
new_left_values,
sharding_func,
name=update_op_name)
# Create the loss subgraph
loss_sp_input = (sparse_ops.sparse_transpose(new_sp_input)
if transpose_input else new_sp_input)
# sp_approx is the low rank estimate of the input matrix, formed by
# computing the product <\\(u_i, v_j\\)> for (i, j) in loss_sp_input.indices.
sp_approx_vals = gen_factorization_ops.masked_matmul(
new_left_values,
right,
loss_sp_input.indices,
transpose_a=False,
transpose_b=True)
sp_approx = sparse_tensor.SparseTensor(loss_sp_input.indices,
sp_approx_vals,
loss_sp_input.dense_shape)
sp_approx_sq = math_ops.square(sp_approx)
sp_residual = sparse_ops.sparse_add(loss_sp_input, sp_approx * (-1))
sp_residual_sq = math_ops.square(sp_residual)
row_wt_mat = (constant_op.constant(0.) if self._row_weights is None
else array_ops.expand_dims(row_weights_slice, 1))
col_wt_mat = (constant_op.constant(0.) if self._col_weights is None
else array_ops.expand_dims(col_weights, 0))
# We return the normalized loss
partial_row_gramian = math_ops.matmul(new_left_values,
new_left_values,
transpose_a=True)
normalization_factor = total_rows / math_ops.cast(
num_rows, dtypes.float32)
unregularized_loss = (
self._unobserved_weight * ( # pyformat line break
sparse_ops.sparse_reduce_sum(sp_residual_sq) - # pyformat break
sparse_ops.sparse_reduce_sum(sp_approx_sq) + # pyformat break
math_ops.trace(math_ops.matmul(partial_row_gramian, gramian)))
+ sparse_ops.sparse_reduce_sum(
row_wt_mat *
(sp_residual_sq * col_wt_mat))) * normalization_factor
if self._regularization is not None:
regularization = self._regularization * (
math_ops.trace(partial_row_gramian) * normalization_factor +
math_ops.trace(gramian))
else:
regularization = constant_op.constant(0.)
sum_weights = self._unobserved_weight * math_ops.cast(
total_rows * total_cols, dtypes.float32)
if self._row_weights is not None and self._col_weights is not None:
ones = sparse_tensor.SparseTensor(
indices=loss_sp_input.indices,
values=array_ops.ones(array_ops.shape(loss_sp_input.values)),
dense_shape=loss_sp_input.dense_shape)
sum_weights += sparse_ops.sparse_reduce_sum(
row_wt_mat * (ones * col_wt_mat)) * normalization_factor
return (new_left_values, update_op, unregularized_loss, regularization,
sum_weights)
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."
" values.shape=%s. weights.shape=%s." %
(_ASSERT_BROADCASTABLE_ERROR_PREFIX, values_rank_static,
weights_rank_static, values.shape, weights.shape))
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 frame(signal,
frame_length,
frame_step,
pad_end=False,
pad_value=0,
axis=-1,
name=None):
"""Expands `signal`'s `axis` dimension into frames of `frame_length`.
Slides a window of size `frame_length` over `signal`'s `axis` dimension
with a stride of `frame_step`, replacing the `axis` dimension with
`[frames, frame_length]` frames.
If `pad_end` is True, window positions that are past the end of the `axis`
dimension are padded with `pad_value` until the window moves fully past the
end of the dimension. Otherwise, only window positions that fully overlap the
`axis` dimension are produced.
For example:
```python
# A batch size 3 tensor of 9152 audio samples.
audio = tf.random.normal([3, 9152])
# Compute overlapping frames of length 512 with a step of 180 (frames overlap
# by 332 samples). By default, only 50 frames are generated since the last
# 152 samples do not form a full frame.
frames = tf.signal.frame(audio, 512, 180)
frames.shape.assert_is_compatible_with([3, 50, 512])
# When pad_end is enabled, the final frame is kept (padded with zeros).
frames = tf.signal.frame(audio, 512, 180, pad_end=True)
frames.shape.assert_is_compatible_with([3, 51, 512])
```
Args:
signal: A `[..., samples, ...]` `Tensor`. The rank and dimensions
may be unknown. Rank must be at least 1.
frame_length: The frame length in samples. An integer or scalar `Tensor`.
frame_step: The frame hop size in samples. An integer or scalar `Tensor`.
pad_end: Whether to pad the end of `signal` with `pad_value`.
pad_value: An optional scalar `Tensor` to use where the input signal
does not exist when `pad_end` is True.
axis: A scalar integer `Tensor` indicating the axis to frame. Defaults to
the last axis. Supports negative values for indexing from the end.
name: An optional name for the operation.
Returns:
A `Tensor` of frames with shape `[..., frames, frame_length, ...]`.
Raises:
ValueError: If `frame_length`, `frame_step`, `pad_value`, or `axis` are not
scalar.
"""
with ops.name_scope(name, "frame",
[signal, frame_length, frame_step, pad_value]):
signal = ops.convert_to_tensor(signal, name="signal")
frame_length = ops.convert_to_tensor(frame_length, name="frame_length")
frame_step = ops.convert_to_tensor(frame_step, name="frame_step")
axis = ops.convert_to_tensor(axis, name="axis")
signal.shape.with_rank_at_least(1)
frame_length.shape.assert_has_rank(0)
frame_step.shape.assert_has_rank(0)
axis.shape.assert_has_rank(0)
result_shape = _infer_frame_shape(signal, frame_length, frame_step,
pad_end, axis)
def maybe_constant(val):
val_static = tensor_util.constant_value(val)
return (val_static, True) if val_static is not None else (val,
False)
signal_shape, signal_shape_is_static = maybe_constant(
array_ops.shape(signal))
axis, axis_is_static = maybe_constant(axis)
if signal_shape_is_static and axis_is_static:
# Axis can be negative. Convert it to positive.
axis = range(len(signal_shape))[axis]
outer_dimensions, length_samples, inner_dimensions = np.split(
signal_shape, indices_or_sections=[axis, axis + 1])
length_samples = length_samples.item()
else:
signal_rank = array_ops.rank(signal)
# Axis can be negative. Convert it to positive.
axis = math_ops.range(signal_rank)[axis]
outer_dimensions, length_samples, inner_dimensions = array_ops.split(
signal_shape, [axis, 1, signal_rank - 1 - axis])
length_samples = array_ops.reshape(length_samples, [])
num_outer_dimensions = array_ops.size(outer_dimensions)
num_inner_dimensions = array_ops.size(inner_dimensions)
# If padding is requested, pad the input signal tensor with pad_value.
if pad_end:
pad_value = ops.convert_to_tensor(pad_value, signal.dtype)
pad_value.shape.assert_has_rank(0)
# Calculate number of frames, using double negatives to round up.
num_frames = -(-length_samples // frame_step)
# Pad the signal by up to frame_length samples based on how many samples
# are remaining starting from last_frame_position.
pad_samples = math_ops.maximum(
0,
frame_length + frame_step * (num_frames - 1) - length_samples)
# Pad the inner dimension of signal by pad_samples.
paddings = array_ops.concat([
array_ops.zeros([num_outer_dimensions, 2],
dtype=pad_samples.dtype), [[0, pad_samples]],
array_ops.zeros([num_inner_dimensions, 2],
dtype=pad_samples.dtype)
], 0)
signal = array_ops.pad(signal, paddings, constant_values=pad_value)
signal_shape = array_ops.shape(signal)
length_samples = signal_shape[axis]
else:
num_frames = math_ops.maximum(
0, 1 + (length_samples - frame_length) // frame_step)
subframe_length, _ = maybe_constant(
util_ops.gcd(frame_length, frame_step))
subframes_per_frame = frame_length // subframe_length
subframes_per_hop = frame_step // subframe_length
num_subframes = length_samples // subframe_length
slice_shape = array_ops.concat([
outer_dimensions, [num_subframes * subframe_length],
inner_dimensions
], 0)
subframe_shape = array_ops.concat([
outer_dimensions, [num_subframes, subframe_length],
inner_dimensions
], 0)
subframes = array_ops.reshape(
array_ops.strided_slice(signal, array_ops.zeros_like(signal_shape),
slice_shape), subframe_shape)
# frame_selector is a [num_frames, subframes_per_frame] tensor
# that indexes into the appropriate frame in subframes. For example:
# [[0, 0, 0, 0], [2, 2, 2, 2], [4, 4, 4, 4]]
frame_selector = array_ops.reshape(
math_ops.range(num_frames) * subframes_per_hop, [num_frames, 1])
# subframe_selector is a [num_frames, subframes_per_frame] tensor
# that indexes into the appropriate subframe within a frame. For example:
# [[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]]
subframe_selector = array_ops.reshape(
math_ops.range(subframes_per_frame), [1, subframes_per_frame])
# Adding the 2 selector tensors together produces a [num_frames,
# subframes_per_frame] tensor of indices to use with tf.gather to select
# subframes from subframes. We then reshape the inner-most
# subframes_per_frame dimension to stitch the subframes together into
# frames. For example: [[0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7]].
selector = frame_selector + subframe_selector
frames = array_ops.reshape(
array_ops.gather(subframes, selector, axis=axis),
array_ops.concat([
outer_dimensions, [num_frames, frame_length], inner_dimensions
], 0))
if result_shape:
frames.set_shape(result_shape)
return frames
def gather_nd(params, indices, batch_dims=0, name=None):
"""Gather slices from `params` using `n`-dimensional indices.
This operation is similar to `gather`, but it uses the innermost dimension
of `indices` to define a slice into `params`. In particular, if:
* `indices` has shape `[A1...AN, I]`
* `params` has shape `[B1...BM]`
Then:
* `result` has shape `[A1...AN, B_{I+1}...BM]`.
* `result[a1...aN] = params[indices[a1...aN, :]]`
Args:
params: A potentially ragged tensor with shape `[A1...AN, I]`.
indices: A potentially ragged tensor with shape `[B1...BM]`.
batch_dims: Must be zero.
name: A name for the operation (optional).
Returns:
A potentially ragged tensor with shape `[A1...AN, B_{I+1}...BM]`.
#### Examples:
```python
>>> params = tf.compat.v1.ragged.constant_value(
... [ [ ['000', '001'], ['010' ] ],
... [ ['100' ], ['110', '111', '112'], ['120'] ],
... [ [ ], ['210' ] ] ])
>>> # Gather 2D slices from a 3D tensor
>>> ragged.gather_nd(params, [[2], [0]])
[ [ [ ], ['210'] ]
[ ['000', '001'], ['010'] ] ]
>>> # Gather 1D slices from a 3D tensor
>>> ragged.gather_nd(params, [[2, 1], [0, 0]])
[['210'], ['000', '001']]
>>> # Gather scalars from a 3D tensor
>>> ragged.gather_nd(params, [[0, 0, 1], [1, 1, 2]])
['001', '112']
```
"""
if not isinstance(batch_dims, int) or batch_dims != 0:
raise ValueError(
'batch_dims != 0 is not supported for ragged gather yet.')
if not (ragged_tensor.is_ragged(params)
or ragged_tensor.is_ragged(indices)):
return array_ops.gather_nd(params, indices, name)
with ops.name_scope(name, 'RaggedGatherNd', [params, indices]):
params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params')
indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
params, indices = ragged_tensor.match_row_splits_dtypes(
params, indices)
indices_shape = indices.shape
indices_ndims = indices_shape.ndims
if indices_ndims is None:
raise ValueError('indices.rank be statically known.')
if indices_ndims == 0:
raise ValueError('indices.rank must be at least 1.')
if (ragged_tensor.is_ragged(indices)
and indices_ndims == indices.ragged_rank + 1):
raise ValueError(
'The innermost dimension of indices may not be ragged')
# `index_size` is the "n" in "gather_nd" -- i.e., the number of dimensions
# that each index slices into.
index_size = tensor_shape.dimension_value(indices_shape[-1])
if index_size is None:
raise ValueError('indices.shape[-1] must be statically known.')
# If `indices` has more than 2 dimensions, then recurse. If `indices` is
# dense, then we convert it to ragged before recursing, and then convert
# the result back to `dense` if appropriate.
if indices_ndims > 2:
indices_is_dense = not ragged_tensor.is_ragged(indices)
if indices_is_dense:
indices = ragged_tensor.RaggedTensor.from_tensor(
indices,
ragged_rank=indices_ndims - 2,
row_splits_dtype=params.row_splits.dtype)
result = indices.with_flat_values(
gather_nd(params, indices.flat_values))
if (indices_is_dense and ragged_tensor.is_ragged(result)
and result.ragged_rank == indices_ndims - 2):
result = ragged_tensor.RaggedTensor.to_tensor(result)
return result
# indices_ndims <= 2, and the innermost dimension of indices may not be
# ragged, so `indices` must not be ragged.
assert not ragged_tensor.is_ragged(indices)
assert ragged_tensor.is_ragged(params)
# Handle corner case: An empty index tuple selects the entire `params`
# value. So if `index_size` is zero, then tile `params`.
if index_size == 0:
params_ndims = params.ragged_rank + array_ops.rank(
params.flat_values)
for dim in range(indices_ndims - 1):
params = ragged_array_ops.expand_dims(params, axis=0)
multiples = array_ops.concat([
array_ops.shape(indices)[:-1],
array_ops.ones([params_ndims], dtypes.int32)
],
axis=0)
return ragged_array_ops.tile(params, multiples)
# When index_size=1, we can just flatten the index tuples and use gather.
elif index_size == 1:
flattened_index_tuples = array_ops.reshape(indices, [-1])
return gather(params, flattened_index_tuples)
# Otherwise, params is a RaggedTensor, and indices is a 1D or 2D Tensor.
# Flatten both the index tuples and the params, such that the flattened
# index tuples point to the correct values in the flattened params; and
# then use ragged.gather on the flattened index tuples & params.
else:
indices = math_ops.cast(indices, params.row_splits.dtype)
# Flatten the outermost 2 dimensions of the index tuples & params.
flattened_index_tuples = array_ops.gather(params.row_splits,
indices[..., 0])
flattened_index_tuples += indices[..., 1]
flattened_params = params.values
# Flatten any remaining dimensions.
for dim in range(2, index_size):
if not ragged_tensor.is_ragged(flattened_params):
flattened_index_tuples = array_ops.expand_dims(
flattened_index_tuples, axis=1)
flattened_index_tuples = array_ops.concat(
[flattened_index_tuples, indices[..., dim:]], axis=1)
return array_ops.gather_nd(flattened_params,
flattened_index_tuples)
flattened_index_tuples = array_ops.gather(
flattened_params.row_starts(), flattened_index_tuples)
flattened_index_tuples += indices[..., dim]
flattened_params = flattened_params.values
# Gather using the flattened index tuples and params.
return gather(flattened_params, flattened_index_tuples)
def target_log_prob(x):
event_dims = math_ops.range(independent_chain_ndims,
array_ops.rank(x))
return self._log_gamma_log_prob(x, event_dims)
def embedding_lookup_sparse(params,
sp_ids,
sp_weights,
partition_strategy="mod",
name=None,
combiner=None,
max_norm=None):
"""Computes embeddings for the given ids and weights.
This op assumes that there is at least one id for each row in the dense tensor
represented by sp_ids (i.e. there are no rows with empty features), and that
all the indices of sp_ids are in canonical row-major order.
It also assumes that all id values lie in the range [0, p0), where p0
is the sum of the size of params along dimension 0.
Args:
params: A single tensor representing the complete embedding tensor,
or a list of P tensors all of same shape except for the first dimension,
representing sharded embedding tensors. Alternatively, a
`PartitionedVariable`, created by partitioning along dimension 0. Each
element must be appropriately sized for the given `partition_strategy`.
sp_ids: N x M SparseTensor of int64 ids (typically from FeatureValueToId),
where N is typically batch size and M is arbitrary.
sp_weights: either a SparseTensor of float / double weights, or None to
indicate all weights should be taken to be 1. If specified, sp_weights
must have exactly the same shape and indices as sp_ids.
partition_strategy: A string specifying the partitioning strategy, relevant
if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
is `"mod"`. See `tf.nn.embedding_lookup` for more details.
name: Optional name for the op.
combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
and "sum" are supported.
"sum" computes the weighted sum of the embedding results for each row.
"mean" is the weighted sum divided by the total weight.
"sqrtn" is the weighted sum divided by the square root of the sum of the
squares of the weights.
max_norm: If provided, each embedding is normalized to have l2 norm equal
to max_norm before combining.
Returns:
A dense tensor representing the combined embeddings for the
sparse ids. For each row in the dense tensor represented by sp_ids, the op
looks up the embeddings for all ids in that row, multiplies them by the
corresponding weight, and combines these embeddings as specified.
In other words, if
shape(combined params) = [p0, p1, ..., pm]
and
shape(sp_ids) = shape(sp_weights) = [d0, d1, ..., dn]
then
shape(output) = [d0, d1, ..., dn-1, p1, ..., pm].
For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are
[0, 0]: id 1, weight 2.0
[0, 1]: id 3, weight 0.5
[1, 0]: id 0, weight 1.0
[2, 3]: id 1, weight 3.0
with `combiner`="mean", then the output will be a 3x20 matrix where
output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
output[1, :] = params[0, :] * 1.0
output[2, :] = params[1, :] * 3.0
Raises:
TypeError: If sp_ids is not a SparseTensor, or if sp_weights is neither
None nor SparseTensor.
ValueError: If combiner is not one of {"mean", "sqrtn", "sum"}.
"""
if combiner is None:
logging.warn("The default value of combiner will change from \"mean\" "
"to \"sqrtn\" after 2016/11/01.")
combiner = "mean"
if combiner not in ("mean", "sqrtn", "sum"):
raise ValueError("combiner must be one of 'mean', 'sqrtn' or 'sum'")
if isinstance(params, variables.PartitionedVariable):
params = list(params) # Iterate to get the underlying Variables.
if not isinstance(params, list):
params = [params]
if not isinstance(sp_ids, sparse_tensor.SparseTensor):
raise TypeError("sp_ids must be SparseTensor")
ignore_weights = sp_weights is None
if not ignore_weights:
if not isinstance(sp_weights, sparse_tensor.SparseTensor):
raise TypeError("sp_weights must be either None or SparseTensor")
sp_ids.values.get_shape().assert_is_compatible_with(
sp_weights.values.get_shape())
sp_ids.indices.get_shape().assert_is_compatible_with(
sp_weights.indices.get_shape())
sp_ids.dense_shape.get_shape().assert_is_compatible_with(
sp_weights.dense_shape.get_shape())
# TODO(yleon): Add enhanced node assertions to verify that sp_ids and
# sp_weights have equal indices and shapes.
with ops.name_scope(name, "embedding_lookup_sparse",
params + [sp_ids]) as name:
segment_ids = sp_ids.indices[:, 0]
if segment_ids.dtype != dtypes.int32:
segment_ids = math_ops.cast(segment_ids, dtypes.int32)
ids = sp_ids.values
if ignore_weights:
ids, idx = array_ops.unique(ids)
else:
idx = None
embeddings = embedding_lookup(params,
ids,
partition_strategy=partition_strategy,
max_norm=max_norm)
if not ignore_weights:
weights = sp_weights.values
if weights.dtype != embeddings.dtype:
weights = math_ops.cast(weights, embeddings.dtype)
# Reshape weights to allow broadcast
ones = array_ops.fill(
array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0), 1)
bcast_weights_shape = array_ops.concat(
[array_ops.shape(weights), ones], 0)
orig_weights_shape = weights.get_shape()
weights = array_ops.reshape(weights, bcast_weights_shape)
# Set the weight shape, since after reshaping to bcast_weights_shape,
# the shape becomes None.
if embeddings.get_shape().ndims is not None:
weights.set_shape(
orig_weights_shape.concatenate(
[1 for _ in range(embeddings.get_shape().ndims - 1)]))
embeddings *= weights
if combiner == "sum":
embeddings = math_ops.segment_sum(embeddings,
segment_ids,
name=name)
elif combiner == "mean":
embeddings = math_ops.segment_sum(embeddings, segment_ids)
weight_sum = math_ops.segment_sum(weights, segment_ids)
embeddings = math_ops.div(embeddings, weight_sum, name=name)
elif combiner == "sqrtn":
embeddings = math_ops.segment_sum(embeddings, segment_ids)
weights_squared = math_ops.pow(weights, 2)
weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
weight_sum_sqrt = math_ops.sqrt(weight_sum)
embeddings = math_ops.div(embeddings,
weight_sum_sqrt,
name=name)
else:
assert False, "Unrecognized combiner"
else:
assert idx is not None
if combiner == "sum":
embeddings = math_ops.sparse_segment_sum(embeddings,
idx,
segment_ids,
name=name)
elif combiner == "mean":
embeddings = math_ops.sparse_segment_mean(embeddings,
idx,
segment_ids,
name=name)
elif combiner == "sqrtn":
embeddings = math_ops.sparse_segment_sqrt_n(embeddings,
idx,
segment_ids,
name=name)
else:
assert False, "Unrecognized combiner"
return embeddings
def mean_pairwise_squared_error(labels,
predictions,
weights=1.0,
scope=None,
loss_collection=ops.GraphKeys.LOSSES):
"""Adds a pairwise-errors-squared loss to the training procedure.
Unlike `mean_squared_error`, which is a measure of the differences between
corresponding elements of `predictions` and `labels`,
`mean_pairwise_squared_error` is a measure of the differences between pairs of
corresponding elements of `predictions` and `labels`.
For example, if `labels`=[a, b, c] and `predictions`=[x, y, z], there are
three pairs of differences are summed to compute the loss:
loss = [ ((a-b) - (x-y)).^2 + ((a-c) - (x-z)).^2 + ((b-c) - (y-z)).^2 ] / 3
Note that since the inputs are of shape `[batch_size, d0, ... dN]`, the
corresponding pairs are computed within each batch sample but not across
samples within a batch. For example, if `predictions` represents a batch of
16 grayscale images of dimension [batch_size, 100, 200], then the set of pairs
is drawn from each image, but not across images.
`weights` acts as a coefficient for the loss. If a scalar is provided, then
the loss is simply scaled by the given value. If `weights` is a tensor of size
[batch_size], then the total loss for each sample of the batch is rescaled
by the corresponding element in the `weights` vector.
Args:
labels: The ground truth output tensor, whose shape must match the shape of
`predictions`.
predictions: The predicted outputs, a tensor of size
`[batch_size, d0, .. dN]` where N+1 is the total number of dimensions in
`predictions`.
weights: Coefficients for the loss a scalar, a tensor of shape
`[batch_size]` or a tensor whose shape matches `predictions`.
scope: The scope for the operations performed in computing the loss.
loss_collection: collection to which the loss will be added.
Returns:
A scalar `Tensor` that returns the weighted loss.
Raises:
ValueError: If the shape of `predictions` doesn't match that of `labels` or
if the shape of `weights` is invalid. Also if `labels` or `predictions`
is None.
"""
if labels is None:
raise ValueError("labels must not be None.")
if predictions is None:
raise ValueError("predictions must not be None.")
with ops.name_scope(scope, "mean_pairwise_squared_error",
(predictions, labels, weights)) as scope:
weights = math_ops.to_float(weights)
labels = math_ops.to_float(labels)
with ops.control_dependencies(
(weights_broadcast_ops.assert_broadcastable(weights, labels), )):
predictions = math_ops.to_float(predictions)
predictions.get_shape().assert_is_compatible_with(
labels.get_shape())
diffs = math_ops.subtract(predictions, labels)
reduction_indices = math_ops.range(1, array_ops.rank(diffs))
sum_squares_diff_per_batch = math_ops.reduce_sum(
math_ops.square(diffs),
reduction_indices=reduction_indices,
keep_dims=True)
num_present_per_batch = _num_present(diffs,
weights,
per_batch=True)
term1 = 2.0 * _safe_div(sum_squares_diff_per_batch,
num_present_per_batch)
sum_diff = math_ops.reduce_sum(diffs,
reduction_indices=reduction_indices,
keep_dims=True)
term2 = 2.0 * _safe_div(math_ops.square(sum_diff),
math_ops.square(num_present_per_batch))
weighted_losses = math_ops.multiply(term1 - term2, weights)
loss = math_ops.reduce_sum(weighted_losses)
mean_loss = array_ops.where(
math_ops.reduce_sum(num_present_per_batch) > 0,
loss,
array_ops.zeros_like(loss),
name="value")
util.add_loss(mean_loss, loss_collection)
return mean_loss
def squeeze_or_expand_dimensions(y_pred, y_true=None, sample_weight=None):
"""Squeeze or expand last dimension if needed.
1. Squeezes last dim of `y_pred` or `y_true` if their rank differs by 1
(using `remove_squeezable_dimensions`).
2. Squeezes or expands last dim of `sample_weight` if its rank differs by 1
from the new rank of `y_pred`.
If `sample_weight` is scalar, it is kept scalar.
This will use static shape if available. Otherwise, it will add graph
operations, which could result in a performance hit.
Args:
y_pred: Predicted values, a `Tensor` of arbitrary dimensions.
y_true: Optional label `Tensor` whose dimensions match `y_pred`.
sample_weight: Optional weight scalar or `Tensor` whose dimensions match
`y_pred`.
Returns:
Tuple of `y_pred`, `y_true` and `sample_weight`. Each of them possibly has
the last dimension squeezed,
`sample_weight` could be extended by one dimension.
If `sample_weight` is None, (y_pred, y_true) is returned.
"""
y_pred_shape = y_pred.shape
y_pred_rank = y_pred_shape.ndims
if y_true is not None:
# If sparse matrix is provided as `y_true`, the last dimension in `y_pred`
# may be > 1. Eg: y_true = [0, 1, 2] (shape=(3,)),
# y_pred = [[.9, .05, .05], [.5, .89, .6], [.05, .01, .94]] (shape=(3, 3))
# In this case, we should not try to remove squeezable dimension.
y_true_shape = y_true.shape
y_true_rank = y_true_shape.ndims
if (y_true_rank is not None) and (y_pred_rank is not None):
# Use static rank for `y_true` and `y_pred`.
if (y_pred_rank - y_true_rank != 1) or y_pred_shape[-1] == 1:
y_true, y_pred = remove_squeezable_dimensions(y_true, y_pred)
else:
# Use dynamic rank.
rank_diff = array_ops.rank(y_pred) - array_ops.rank(y_true)
squeeze_dims = lambda: remove_squeezable_dimensions( # pylint: disable=g-long-lambda
y_true, y_pred)
is_last_dim_1 = math_ops.equal(1, array_ops.shape(y_pred)[-1])
maybe_squeeze_dims = lambda: control_flow_ops.cond( # pylint: disable=g-long-lambda
is_last_dim_1, squeeze_dims, lambda: (y_true, y_pred))
y_true, y_pred = control_flow_ops.cond(
math_ops.equal(1, rank_diff), maybe_squeeze_dims, squeeze_dims)
if sample_weight is None:
return y_pred, y_true
weights_shape = sample_weight.shape
weights_rank = weights_shape.ndims
if weights_rank == 0: # If weights is scalar, do nothing.
return y_pred, y_true, sample_weight
if (y_pred_rank is not None) and (weights_rank is not None):
# Use static rank.
if weights_rank - y_pred_rank == 1:
sample_weight = array_ops.squeeze(sample_weight, [-1])
elif y_pred_rank - weights_rank == 1:
sample_weight = array_ops.expand_dims(sample_weight, [-1])
return y_pred, y_true, sample_weight
# Use dynamic rank.
weights_rank_tensor = array_ops.rank(sample_weight)
rank_diff = weights_rank_tensor - array_ops.rank(y_pred)
maybe_squeeze_weights = lambda: array_ops.squeeze(sample_weight, [-1])
def _maybe_expand_weights():
expand_weights = lambda: array_ops.expand_dims(sample_weight, [-1])
return control_flow_ops.cond(math_ops.equal(rank_diff, -1),
expand_weights, lambda: sample_weight)
def _maybe_adjust_weights():
return control_flow_ops.cond(math_ops.equal(rank_diff,
1), maybe_squeeze_weights,
_maybe_expand_weights)
# squeeze or expand last dim of `sample_weight` if its rank differs by 1
# from the new rank of `y_pred`.
sample_weight = control_flow_ops.cond(
math_ops.equal(weights_rank_tensor, 0), lambda: sample_weight,
_maybe_adjust_weights)
return y_pred, y_true, sample_weight
def _ReductionDims(x, reduction_indices):
"""Returns range(0, rank(x)) if reduction_indices is None."""
if reduction_indices is not None:
return reduction_indices
else:
return range(0, array_ops.rank(x))
def _validate_sample_arg(self, x):
"""Helper which validates sample arg, e.g., input to `log_prob`."""
with ops.name_scope(name="validate_sample_arg", values=[x]):
x_ndims = (array_ops.rank(x)
if x.shape.ndims is None else x.shape.ndims)
event_ndims = (array_ops.size(self.event_shape_tensor())
if self.event_shape.ndims is None else
self.event_shape.ndims)
batch_ndims = (array_ops.size(self._batch_shape_unexpanded)
if self.batch_shape.ndims is None else
self.batch_shape.ndims)
expected_batch_event_ndims = batch_ndims + event_ndims
if (isinstance(x_ndims, int)
and isinstance(expected_batch_event_ndims, int)):
if x_ndims < expected_batch_event_ndims:
raise NotImplementedError(
"Broadcasting is not supported; too few batch and event dims "
"(expected at least {}, saw {}).".format(
expected_batch_event_ndims, x_ndims))
ndims_assertion = []
elif self.validate_args:
ndims_assertion = [
check_ops.assert_greater_equal(
x_ndims,
expected_batch_event_ndims,
message=("Broadcasting is not supported; too few "
"batch and event dims."),
name="assert_batch_and_event_ndims_large_enough"),
]
if (self.batch_shape.is_fully_defined()
and self.event_shape.is_fully_defined()):
expected_batch_event_shape = np.int32(
self.batch_shape.concatenate(self.event_shape).as_list())
else:
expected_batch_event_shape = array_ops.concat([
self.batch_shape_tensor(),
self.event_shape_tensor(),
],
axis=0)
sample_ndims = x_ndims - expected_batch_event_ndims
if isinstance(sample_ndims, int):
sample_ndims = max(sample_ndims, 0)
if (isinstance(sample_ndims, int)
and x.shape[sample_ndims:].is_fully_defined()):
actual_batch_event_shape = np.int32(
x.shape[sample_ndims:].as_list())
else:
sample_ndims = math_ops.maximum(sample_ndims, 0)
actual_batch_event_shape = array_ops.shape(x)[sample_ndims:]
if (isinstance(expected_batch_event_shape, np.ndarray)
and isinstance(actual_batch_event_shape, np.ndarray)):
if any(expected_batch_event_shape != actual_batch_event_shape):
raise NotImplementedError(
"Broadcasting is not supported; "
"unexpected batch and event shape "
"(expected {}, saw {}).".format(
expected_batch_event_shape,
actual_batch_event_shape))
# We need to set the final runtime-assertions to `ndims_assertion` since
# its possible this assertion was created. We could add a condition to
# only do so if `self.validate_args == True`, however this is redundant
# as `ndims_assertion` already encodes this information.
runtime_assertions = ndims_assertion
elif self.validate_args:
# We need to make the `ndims_assertion` a control dep because otherwise
# TF itself might raise an exception owing to this assertion being
# ill-defined, ie, one cannot even compare different rank Tensors.
with ops.control_dependencies(ndims_assertion):
shape_assertion = check_ops.assert_equal(
expected_batch_event_shape,
actual_batch_event_shape,
message=("Broadcasting is not supported; "
"unexpected batch and event shape."),
name="assert_batch_and_event_shape_same")
runtime_assertions = [shape_assertion]
else:
runtime_assertions = []
return runtime_assertions
def embedding_lookup_sparse(
params,
sp_ids,
sp_weights,
partition_strategy=None, # no used
name="embedding_lookup_sparse",
combiner="mean",
max_norm=None,
return_trainable=False,
):
"""Provides a dynamic version of embedding_lookup_sparse
similar with tf.nn.embedding_lookup_sparse.
This op assumes that there is at least one id for each row in the dense tensor
represented by sp_ids (i.e. there are no rows with empty features), and that
all the indices of sp_ids are in canonical row-major order.
It also assumes that all id values lie in the range [0, p0), where p0
is the sum of the size of params along dimension 0.
Args:
params: A single `dynamic_embedding.Variable` instance representing
the complete embedding tensor.
sp_ids: N x M `SparseTensor` of int64 ids where N is typically batch size
and M is arbitrary.
sp_weights: either a `SparseTensor` of float / double weights, or `None` to
indicate all weights should be taken to be 1. If specified, `sp_weights`
must have exactly the same shape and indices as `sp_ids`.
partition_strategy: No used.
name: Optional name for the op.
combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
and "sum" are supported. "sum" computes the weighted sum of the embedding
results for each row. "mean" is the weighted sum divided by the total
weight. "sqrtn" is the weighted sum divided by the square root of the sum
of the squares of the weights.
max_norm: If not `None`, each embedding is clipped if its l2-norm is larger
than this value, before combining.
return_trainable: optional, If True, also return TrainableWrapper create by
`dynamic_embedding.embedding_lookup`
Returns:
combined_embeddings: A dense tensor representing the combined embeddings
for the sparse ids. For each row in the dense tensor represented by
`sp_ids`, the op looks up the embeddings for all ids in that row,
multiplies them by the corresponding weight, and combines these embeddings
as specified.
In other words, if
`shape(combined params) = [+infinity, dim]`
and
`shape(sp_ids) = shape(sp_weights) = [d0, d1, ..., dn]`
then
`shape(output) = [d0, dim]`.
For instance, if params dim=20, and sp_ids / sp_weights are
```python
[0, 0]: id 1, weight 2.0
[0, 1]: id 3, weight 0.5
[1, 0]: id 0, weight 1.0
[2, 3]: id 1, weight 3.0
```
with `combiner`="mean", then the output will be a 3x20 matrix where
```python
output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
output[1, :] = (params[0, :] * 1.0) / 1.0
output[2, :] = (params[1, :] * 3.0) / 3.0
```
trainable_wrap:
A TrainableWrapper object used to fill the Optimizers `var_list`
Only provided if `return_trainable` is True.
Raises:
TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is
neither `None` nor `SparseTensor`.
ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}.
"""
if combiner not in ("mean", "sqrtn", "sum"):
raise ValueError("combiner must be one of 'mean', 'sqrtn' or 'sum'")
if not isinstance(sp_ids, sparse_tensor.SparseTensor):
raise TypeError("sp_ids must be SparseTensor")
ignore_weights = sp_weights is None
if not ignore_weights:
if not isinstance(sp_weights, sparse_tensor.SparseTensor):
raise TypeError("sp_weights must be either None or SparseTensor")
scope = variable_scope.get_variable_scope()
full_name = scope.name + "/" + name if scope.name else name
with ops.name_scope(full_name + "/"):
segment_ids = sp_ids.indices[:, 0]
if segment_ids.dtype != dtypes.int32:
segment_ids = math_ops.cast(segment_ids, dtypes.int32)
ids = sp_ids.values
ids, idx = array_ops.unique(ids)
embeddings, trainable_ = embedding_lookup(
params,
ids,
name=name + "/embedding_lookup",
partition_strategy=partition_strategy,
max_norm=max_norm,
return_trainable=True,
)
if embeddings.dtype in (dtypes.float16, dtypes.bfloat16):
embeddings = math_ops.cast(embeddings, dtypes.float32)
if not ignore_weights:
weights = sp_weights.values
if weights.dtype != embeddings.dtype:
weights = math_ops.cast(weights, embeddings.dtype)
embeddings = array_ops.gather(embeddings, idx)
# Reshape weights to allow broadcast
ones = array_ops.fill(
array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0), 1)
bcast_weights_shape = array_ops.concat(
[array_ops.shape(weights), ones], 0)
orig_weights_shape = weights.get_shape()
weights = array_ops.reshape(weights, bcast_weights_shape)
# Set the weight shape, since after reshaping to bcast_weights_shape,
# the shape becomes None.
if embeddings.get_shape().ndims is not None:
weights.set_shape(
orig_weights_shape.concatenate(
[1 for _ in range(embeddings.get_shape().ndims - 1)]))
embeddings *= weights
if combiner == "sum":
embeddings = math_ops.segment_sum(embeddings,
segment_ids,
name=name)
elif combiner == "mean":
embeddings = math_ops.segment_sum(embeddings, segment_ids)
weight_sum = math_ops.segment_sum(weights, segment_ids)
embeddings = math_ops.div(embeddings, weight_sum, name=name)
elif combiner == "sqrtn":
embeddings = math_ops.segment_sum(embeddings, segment_ids)
weights_squared = math_ops.pow(weights, 2)
weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
weight_sum_sqrt = math_ops.sqrt(weight_sum)
embeddings = math_ops.div(embeddings,
weight_sum_sqrt,
name=name)
else:
assert False, "Unrecognized combiner"
else:
assert idx is not None
if combiner == "sum":
embeddings = math_ops.sparse_segment_sum(embeddings,
idx,
segment_ids,
name=name)
elif combiner == "mean":
embeddings = math_ops.sparse_segment_mean(embeddings,
idx,
segment_ids,
name=name)
elif combiner == "sqrtn":
embeddings = math_ops.sparse_segment_sqrt_n(embeddings,
idx,
segment_ids,
name=name)
else:
assert False, "Unrecognized combiner"
return (embeddings, trainable_) if return_trainable else embeddings
def posterior_from_prior_state(self, prior_state, prior_state_var,
observation, observation_model,
predicted_observations, observation_noise):
"""Compute a posterior over states given an observation.
Args:
prior_state: Prior state mean [batch size x state dimension]
prior_state_var: Prior state covariance [batch size x state dimension x
state dimension]
observation: The observed value corresponding to the predictions given
[batch size x observation dimension]
observation_model: The [batch size x observation dimension x model state
dimension] Tensor indicating how a particular state is mapped to
(pre-noise) observations for each part of the batch.
predicted_observations: An (observation mean, observation variance) tuple
computed based on the current state, usually the output of
observed_from_state.
observation_noise: A [batch size x observation dimension x observation
dimension] or [observation dimension x observation dimension] Tensor
with covariance matrices to use for each part of the batch (a
two-dimensional input will be broadcast).
Returns:
Posterior mean and covariance (dimensions matching the first two
arguments).
"""
observed_mean, observed_var = predicted_observations
residual = observation - observed_mean
# TODO(allenl): Can more of this be done using matrix_solve_ls?
kalman_solve_rhs = math_ops.matmul(observation_model,
prior_state_var,
adjoint_b=True)
# This matrix_solve adjoint doesn't make a difference symbolically (since
# observed_var is a covariance matrix, and should be symmetric), but
# filtering on multivariate series is unstable without it. See
# test_multivariate_symmetric_covariance_float64 in kalman_filter_test.py
# for an example of the instability (fails with adjoint=False).
kalman_gain_transposed = linalg_ops.matrix_solve(matrix=observed_var,
rhs=kalman_solve_rhs,
adjoint=True)
posterior_state = prior_state + array_ops.squeeze(math_ops.matmul(
kalman_gain_transposed,
array_ops.expand_dims(residual, -1),
adjoint_a=True),
squeeze_dims=[-1])
gain_obs = math_ops.matmul(kalman_gain_transposed,
observation_model,
adjoint_a=True)
identity_extradim = linalg_ops.eye(array_ops.shape(gain_obs)[1],
dtype=gain_obs.dtype)[None]
identity_minus_factor = identity_extradim - gain_obs
if self._simplified_posterior_covariance_computation:
# posterior covariance =
# (I - kalman_gain * observation_model) * prior_state_var
posterior_state_var = math_ops.matmul(identity_minus_factor,
prior_state_var)
else:
observation_noise = ops.convert_to_tensor(observation_noise)
# A Joseph form update, which provides better numeric stability than the
# simplified optimal Kalman gain update, at the cost of a few extra
# operations. Joseph form updates are valid for any gain (not just the
# optimal Kalman gain), and so are more forgiving of numerical errors in
# computing the optimal Kalman gain.
#
# posterior covariance =
# (I - kalman_gain * observation_model) * prior_state_var
# * (I - kalman_gain * observation_model)^T
# + kalman_gain * observation_noise * kalman_gain^T
left_multiplied_state_var = math_ops.matmul(
identity_minus_factor, prior_state_var)
multiplied_state_var = math_ops.matmul(identity_minus_factor,
left_multiplied_state_var,
adjoint_b=True)
def _batch_observation_noise_update():
return (multiplied_state_var + math_ops.matmul(
math_ops.matmul(kalman_gain_transposed,
observation_noise,
adjoint_a=True), kalman_gain_transposed))
def _matrix_observation_noise_update():
return (multiplied_state_var + math_ops.matmul(
math_utils.batch_times_matrix(
kalman_gain_transposed, observation_noise, adj_x=True),
kalman_gain_transposed))
if observation_noise.get_shape().ndims is None:
posterior_state_var = control_flow_ops.cond(
math_ops.equal(array_ops.rank(observation_noise),
2), _matrix_observation_noise_update,
_batch_observation_noise_update)
else:
# If static shape information exists, it gets checked in each cond()
# branch, so we need a special case to avoid graph-build-time
# exceptions.
if observation_noise.get_shape().ndims == 2:
posterior_state_var = _matrix_observation_noise_update()
else:
posterior_state_var = _batch_observation_noise_update()
return posterior_state, posterior_state_var
def _rank(x):
rank = ops.convert_to_tensor(x).get_shape().ndims
if rank:
return rank, True
else:
return array_ops.rank(x), False
def inverse_stft(stfts,
frame_length,
frame_step,
fft_length=None,
window_fn=window_ops.hann_window,
name=None):
"""Computes the inverse [Short-time Fourier Transform][stft] of `stfts`.
To reconstruct an original waveform, a complimentary window function should
be used in inverse_stft. Such a window function can be constructed with
tf.signal.inverse_stft_window_fn.
Example:
```python
frame_length = 400
frame_step = 160
waveform = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1000])
stft = tf.signal.stft(waveform, frame_length, frame_step)
inverse_stft = tf.signal.inverse_stft(
stft, frame_length, frame_step,
window_fn=tf.signal.inverse_stft_window_fn(frame_step))
```
if a custom window_fn is used in stft, it must be passed to
inverse_stft_window_fn:
```python
frame_length = 400
frame_step = 160
window_fn = functools.partial(window_ops.hamming_window, periodic=True),
waveform = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1000])
stft = tf.signal.stft(
waveform, frame_length, frame_step, window_fn=window_fn)
inverse_stft = tf.signal.inverse_stft(
stft, frame_length, frame_step,
window_fn=tf.signal.inverse_stft_window_fn(
frame_step, forward_window_fn=window_fn))
```
Implemented with GPU-compatible ops and supports gradients.
Args:
stfts: A `complex64` `[..., frames, fft_unique_bins]` `Tensor` of STFT bins
representing a batch of `fft_length`-point STFTs where `fft_unique_bins`
is `fft_length // 2 + 1`
frame_length: An integer scalar `Tensor`. The window length in samples.
frame_step: An integer scalar `Tensor`. The number of samples to step.
fft_length: An integer scalar `Tensor`. The size of the FFT that produced
`stfts`. If not provided, uses the smallest power of 2 enclosing
`frame_length`.
window_fn: A callable that takes a window length and a `dtype` keyword
argument and returns a `[window_length]` `Tensor` of samples in the
provided datatype. If set to `None`, no windowing is used.
name: An optional name for the operation.
Returns:
A `[..., samples]` `Tensor` of `float32` signals representing the inverse
STFT for each input STFT in `stfts`.
Raises:
ValueError: If `stfts` is not at least rank 2, `frame_length` is not scalar,
`frame_step` is not scalar, or `fft_length` is not scalar.
[stft]: https://en.wikipedia.org/wiki/Short-time_Fourier_transform
"""
with ops.name_scope(name, 'inverse_stft', [stfts]):
stfts = ops.convert_to_tensor(stfts, name='stfts')
stfts.shape.with_rank_at_least(2)
frame_length = ops.convert_to_tensor(frame_length, name='frame_length')
frame_length.shape.assert_has_rank(0)
frame_step = ops.convert_to_tensor(frame_step, name='frame_step')
frame_step.shape.assert_has_rank(0)
if fft_length is None:
fft_length = _enclosing_power_of_two(frame_length)
else:
fft_length = ops.convert_to_tensor(fft_length, name='fft_length')
fft_length.shape.assert_has_rank(0)
real_frames = fft_ops.irfft(stfts, [fft_length])
# frame_length may be larger or smaller than fft_length, so we pad or
# truncate real_frames to frame_length.
frame_length_static = tensor_util.constant_value(frame_length)
# If we don't know the shape of real_frames's inner dimension, pad and
# truncate to frame_length.
if (frame_length_static is None or real_frames.shape.ndims is None
or real_frames.shape[-1].value is None):
real_frames = real_frames[..., :frame_length]
real_frames_rank = array_ops.rank(real_frames)
real_frames_shape = array_ops.shape(real_frames)
paddings = array_ops.concat([
array_ops.zeros([real_frames_rank - 1, 2],
dtype=frame_length.dtype),
[[
0,
math_ops.maximum(0, frame_length - real_frames_shape[-1])
]]
], 0)
real_frames = array_ops.pad(real_frames, paddings)
# We know real_frames's last dimension and frame_length statically. If they
# are different, then pad or truncate real_frames to frame_length.
elif real_frames.shape[-1].value > frame_length_static:
real_frames = real_frames[..., :frame_length_static]
elif real_frames.shape[-1].value < frame_length_static:
pad_amount = frame_length_static - real_frames.shape[-1].value
real_frames = array_ops.pad(
real_frames,
[[0, 0]] * (real_frames.shape.ndims - 1) + [[0, pad_amount]])
# The above code pads the inner dimension of real_frames to frame_length,
# but it does so in a way that may not be shape-inference friendly.
# Restore shape information if we are able to.
if frame_length_static is not None and real_frames.shape.ndims is not None:
real_frames.set_shape([None] * (real_frames.shape.ndims - 1) +
[frame_length_static])
# Optionally window and overlap-add the inner 2 dimensions of real_frames
# into a single [samples] dimension.
if window_fn is not None:
window = window_fn(frame_length, dtype=stfts.dtype.real_dtype)
real_frames *= window
return reconstruction_ops.overlap_and_add(real_frames, frame_step)
def log_gamma_log_prob(x):
counter["target_calls"] += 1
event_dims = math_ops.range(independent_chain_ndims,
array_ops.rank(x))
return self._log_gamma_log_prob(x, event_dims)
