def range(start, limit, delta=1, name="range"):
"""Creates a sequence of integers.
This operation creates a sequence of integers that begins at `start` and
extends by increments of `delta` up to but not including `limit`.
For example:
```
# 'start' is 3
# 'limit' is 18
# 'delta' is 3
tf.range(start, limit, delta) ==> [3, 6, 9, 12, 15]
```
Args:
start: A 0-D (scalar) of type `int32`. First entry in sequence.
limit: A 0-D (scalar) of type `int32`. Upper limit of sequence,
exclusive.
delta: A 0-D `Tensor` (scalar) of type `int32`. Optional. Default is 1.
Number that increments `start`.
name: A name for the operation (optional).
Returns:
An 1-D `int32` `Tensor`.
"""
return gen_math_ops._range(start, limit, delta, name=name)
def range(start, limit, delta=1, name="range"):
"""Creates a sequence of integers.
This operation creates a sequence of integers that begins at `start` and
extends by increments of `delta` up to but not including `limit`.
For example:
```
# 'start' is 3
# 'limit' is 18
# 'delta' is 3
tf.range(start, limit, delta) ==> [3, 6, 9, 12, 15]
```
Args:
start: A 0-D (scalar) of type `int32`. First entry in sequence.
limit: A 0-D (scalar) of type `int32`. Upper limit of sequence,
exclusive.
delta: A 0-D `Tensor` (scalar) of type `int32`. Optional. Default is 1.
Number that increments `start`.
name: A name for the operation (optional).
Returns:
An 1-D `int32` `Tensor`.
"""
return gen_math_ops._range(start, limit, delta, name=name)
def call(self, inputs, state):
add = math_ops.add
sub = math_ops.subtract
mult = math_ops.multiply
# computing m_t
m_t = add(math_ops.matmul(state, self._u),
math_ops.matmul(inputs, self._v))
m_t = nn_ops.bias_add(m_t, self._b)
m_t = math_ops.sigmoid(m_t)
# add L1 loss
ops.add_to_collection('L1 loss', math_ops.abs(m_t - self._m_target))
# computing e_t (= thr)
i = gen_math_ops._range(1, self._num_units + 1, 1)
i = math_ops.cast(i, dtype=dtypes.float32)
mtD = gen_array_ops.tile(mult(m_t[1], self._num_units),
[self._num_units])
thr = math_ops.sigmoid(mult(self._sharpness, sub(mtD, i)))
thr = math_ops.round(add(thr, sub(0.5, self._epsilon)))
ones = array_ops.ones_like(thr)
thr_inv = sub(ones, thr)
# computing h_t
gate_inputs = math_ops.matmul(array_ops.concat([inputs, state], 1),
self._kernel)
gate_inputs = nn_ops.bias_add(gate_inputs, self._bias)
output = self._activation(gate_inputs)
output = add(mult(gate_inputs, thr), mult(state, thr_inv))
return output, output
def batch_gather(params, indices, name=None):
"""Gather slices from `params` according to `indices` with leading batch dims.
This operation assumes that the leading dimensions of `indices` are dense,
and the gathers on the axis corresponding to the last dimension of `indices`.
More concretely it computes:
result[i1, ..., in] = params[i1, ..., in-1, indices[i1, ..., in]]
Therefore `params` should be a Tensor of shape [A1, ..., AN, B1, ..., BM],
`indices` should be a Tensor of shape [A1, ..., AN-1, C] and `result` will be
a Tensor of size `[A1, ..., AN-1, C, B1, ..., BM]`.
In the case in which indices is a 1D tensor, this operation is equivalent to
`tf.gather`.
See also `tf.gather` and `tf.gather_nd`.
Args:
params: A Tensor. The tensor from which to gather values.
indices: A Tensor. Must be one of the following types: int32, int64. Index
tensor. Must be in range `[0, params.shape[axis]`, where `axis` is the
last dimension of `indices` itself.
name: A name for the operation (optional).
Returns:
A Tensor. Has the same type as `params`.
Raises:
ValueError: if `indices` has an unknown shape.
"""
with ops.name_scope(name):
indices = ops.convert_to_tensor(indices, name="indices")
params = ops.convert_to_tensor(params, name="params")
indices_shape = tf.shape(indices)
params_shape = tf.shape(params)
ndims = indices.shape.ndims
if ndims is None:
raise ValueError("batch_gather does not allow indices with unknown "
"shape.")
batch_indices = indices
accum_dim_value = 1
for dim in range(ndims-1, 0, -1):
dim_value = params_shape[dim-1]
accum_dim_value *= params_shape[dim]
dim_indices = gen_math_ops._range(0, dim_value, 1)
dim_indices *= accum_dim_value
dim_shape = tf.stack([1] * (dim - 1) + [dim_value] + [1] * (ndims - dim),
axis=0)
batch_indices += tf.cast(tf.reshape(dim_indices, dim_shape), tf.int64)
flat_indices = tf.reshape(batch_indices, [-1])
outer_shape = params_shape[ndims:]
flat_inner_shape = gen_math_ops.prod(
params_shape[:ndims], [0], False)
flat_params = tf.reshape(
params, tf.concat([[flat_inner_shape], outer_shape], axis=0))
flat_result = tf.gather(flat_params, flat_indices)
result = tf.reshape(flat_result, tf.concat([indices_shape, outer_shape], axis=0))
final_shape = indices.get_shape()[:ndims-1].merge_with(
params.get_shape()[:ndims -1])
final_shape = final_shape.concatenate(indices.get_shape()[ndims-1])
final_shape = final_shape.concatenate(params.get_shape()[ndims:])
result.set_shape(final_shape)
return result
def transpose(a, perm=None, name="transpose"):
"""Transposes `a`. Permutes the dimensions according to `perm`.
The returned tensor's dimension i will correspond to the input dimension
`perm[i]`. If `perm` is not given, it is set to (n-1...0), where n is
the rank of the input tensor. Hence by default, this operation performs a
regular matrix transpose on 2-D input Tensors.
For example:
```python
# 'x' is [[1 2 3]
# [4 5 6]]
tf.transpose(x) ==> [[1 4]
[2 5]
[3 6]]
# Equivalently
tf.transpose(x, perm=[1, 0]) ==> [[1 4]
[2 5]
[3 6]]
# 'perm' is more useful for n-dimensional tensors, for n > 2
# 'x' is [[[1 2 3]
# [4 5 6]]
# [[7 8 9]
# [10 11 12]]]
# Take the transpose of the matrices in dimension-0
tf.transpose(b, perm=[0, 2, 1]) ==> [[[1 4]
[2 5]
[3 6]]
[[7 10]
[8 11]
[9 12]]]
```
Args:
a: A `Tensor`.
perm: A permutation of the dimensions of `a`.
name: A name for the operation (optional).
Returns:
A transposed `Tensor`.
"""
with ops.op_scope([a], name, "transpose") as name:
if perm is None:
dims = gen_math_ops._range(0, gen_array_ops.rank(a), 1)
perm = gen_array_ops.reverse(dims, [True])
ret = gen_array_ops.transpose(a, perm, name=name)
# NOTE(mrry): Setting the shape explicitly because
# reverse is not handled by the shape function.
input_shape = ret.op.inputs[0].get_shape().dims
if input_shape is not None:
ret.set_shape(input_shape[::-1])
else:
ret = gen_array_ops.transpose(a, perm, name=name)
return ret
def transpose(a, perm=None, name="transpose"):
"""Transposes `a`. Permutes the dimensions according to `perm`.
The returned tensor's dimension i will correspond to the input dimension
`perm[i]`. If `perm` is not given, it is set to (n-1...0), where n is
the rank of the input tensor. Hence by default, this operation performs a
regular matrix transpose on 2-D input Tensors.
For example:
```python
# 'x' is [[1 2 3]
# [4 5 6]]
tf.transpose(x) ==> [[1 4]
[2 5]
[3 6]]
# Equivalently
tf.transpose(x, perm=[1, 0]) ==> [[1 4]
[2 5]
[3 6]]
# 'perm' is more useful for n-dimensional tensors, for n > 2
# 'x' is [[[1 2 3]
# [4 5 6]]
# [[7 8 9]
# [10 11 12]]]
# Take the transpose of the matrices in dimension-0
tf.transpose(b, perm=[0, 2, 1]) ==> [[[1 4]
[2 5]
[3 6]]
[[7 10]
[8 11]
[9 12]]]
```
Args:
a: A `Tensor`.
perm: A permutation of the dimensions of `a`.
name: A name for the operation (optional).
Returns:
A transposed `Tensor`.
"""
with ops.op_scope([a], name, "transpose") as name:
if perm is None:
dims = gen_math_ops._range(0, gen_array_ops.rank(a), 1)
perm = gen_array_ops.reverse(dims, [True])
ret = gen_array_ops.transpose(a, perm, name=name)
# NOTE(mrry): Setting the shape explicitly because
# reverse is not handled by the shape function.
input_shape = ret.op.inputs[0].get_shape().dims
if input_shape is not None:
ret.set_shape(input_shape[::-1])
else:
ret = gen_array_ops.transpose(a, perm, name=name)
return ret
def sequence_mask_mid(lengths, maxlen=None, dtype=dtypes.bool, name=None):
"""Return a mask tensor representing the first N positions of each row.
Example:
```python
tf.sequence_mask([1, 3, 2], 5) =
[[True, False, False, False, False],
[False,False, False, False, False],
[False, True, False, False, False]]
```
Args:
lengths: 1D integer tensor, all its values < maxlen.
maxlen: scalar integer tensor, maximum length of each row. Default: use
maximum over lengths.
dtype: output type of the resulting tensor.
name: name of the op.
Returns:
A 2D mask tensor, as shown in the example above, cast to specified dtype.
Raises:
ValueError: if the arguments have invalid rank.
"""
with ops.name_scope(name, "SequenceMask", [lengths, maxlen]):
lengths = (np.array(lengths)) / 2
lengths = ops.convert_to_tensor(lengths, dtype=tf.int32)
# lengths = ops.convert_to_tensor(lengths)
if lengths.get_shape().ndims != 1:
raise ValueError("lengths must be 1D for sequence_mask")
if maxlen is None:
maxlen = gen_math_ops._max(lengths, [0])
else:
maxlen = ops.convert_to_tensor(maxlen)
if maxlen.get_shape().ndims != 0:
raise ValueError("maxlen must be scalar for sequence_mask")
# The basic idea is to compare a range row vector of size maxlen:
# [0, 1, 2, 3, 4]
# to length as a matrix with 1 column: [[1], [3], [2]].
# Because of broadcasting on both arguments this comparison results
# in a matrix of size (len(lengths), maxlen)
row_vector = gen_math_ops._range(constant(0, maxlen.dtype), maxlen,
constant(1, maxlen.dtype))
# Since maxlen >= max(lengths), it is safe to use maxlen as a cast
# authoritative type. Whenever maxlen fits into tf.int32, so do the lengths.
matrix_0 = gen_math_ops.cast(expand_dims(lengths, 1), maxlen.dtype)
matrix_1 = gen_math_ops.cast(expand_dims(lengths - 1, 1), maxlen.dtype)
result_0 = (row_vector < matrix_0)
result_1 = (row_vector >= matrix_1)
result = tf.logical_and(result_0, result_1)
if dtype is None or result.dtype.base_dtype == dtype.base_dtype:
return result
else:
return gen_math_ops.cast(result, dtype)
def sequence_mask(lengths, maxlen=None, dtype=dtypes.bool, name=None):
"""Returns a mask tensor representing the first N positions of each cell.
If `lengths` has shape `[d_1, d_2, ..., d_n]` the resulting tensor `mask` has
dtype `dtype` and shape `[d_1, d_2, ..., d_n, maxlen]`, with
```
mask[i_1, i_2, ..., i_n, j] = (j < lengths[i_1, i_2, ..., i_n])
```
Examples:
```python
tf.sequence_mask([1, 3, 2], 5) # [[True, False, False, False, False],
# [True, True, True, False, False],
# [True, True, False, False, False]]
tf.sequence_mask([[1, 3],[2,0]]) # [[[True, False, False],
# [True, True, True]],
# [[True, True, False],
# [False, False, False]]]
```
Args:
lengths: integer tensor, all its values <= maxlen.
maxlen: scalar integer tensor, size of last dimension of returned tensor.
Default is the maximum value in `lengths`.
dtype: output type of the resulting tensor.
name: name of the op.
Returns:
A mask tensor of shape `lengths.shape + (maxlen,)`, cast to specified dtype.
Raises:
ValueError: if `maxlen` is not a scalar.
"""
with ops.name_scope(name, "SequenceMask", [lengths, maxlen]):
lengths = ops.convert_to_tensor(lengths)
if maxlen is None:
maxlen = gen_math_ops._max(lengths, _all_dimensions(lengths))
maxlen = gen_math_ops.maximum(constant(0, maxlen.dtype), maxlen)
else:
maxlen = ops.convert_to_tensor(maxlen)
if maxlen.get_shape(
).ndims is not None and maxlen.get_shape().ndims != 0:
raise ValueError("maxlen must be scalar for sequence_mask")
# The basic idea is to compare a range row vector of size maxlen:
# [0, 1, 2, 3, 4]
# to length as a matrix with 1 column: [[1], [3], [2]].
# Because of broadcasting on both arguments this comparison results
# in a matrix of size (len(lengths), maxlen)
row_vector = gen_math_ops._range(constant(0, maxlen.dtype), maxlen,
constant(1, maxlen.dtype))
# Since maxlen >= max(lengths), it is safe to use maxlen as a cast
# authoritative type. Whenever maxlen fits into tf.int32, so do the lengths.
matrix = gen_math_ops.cast(expand_dims(lengths, -1), maxlen.dtype)
result = row_vector < matrix
if dtype is None or result.dtype.base_dtype == dtype.base_dtype:
return result
else:
return gen_math_ops.cast(result, dtype)
def _all_dimensions(x):
"""Returns a 1D-tensor listing all dimensions in x."""
# Fast path: avoid creating Rank and Range ops if ndims is known.
if isinstance(x, ops.Tensor) and x.get_shape().ndims is not None:
return constant_op.constant(
np.arange(x.get_shape().ndims), dtype=dtypes.int32)
if (isinstance(x, sparse_tensor.SparseTensor) and
x.dense_shape.get_shape().is_fully_defined()):
r = x.dense_shape.get_shape().dims[0].value # sparse.dense_shape is 1-D.
return constant_op.constant(np.arange(r), dtype=dtypes.int32)
# Otherwise, we rely on `range` and `rank` to do the right thing at runtime.
return gen_math_ops._range(0, rank(x), 1)
def transpose(a, perm=None, name="transpose", conjugate=False):
with ops.name_scope(name, "transpose", [a]) as name:
transpose_fn = (gen_array_ops.conjugate_transpose if
(conjugate
and a.dtype.is_complex) else gen_array_ops.transpose)
if perm is None:
rank = gen_array_ops.rank(a)
perm = (rank - 1) - gen_math_ops._range(0, rank, 1)
ret = transpose_fn(a, perm, name=name)
if not context.executing_eagerly():
input_shape = ret.op.inputs[0].get_shape().dims
if input_shape is not None:
ret.set_shape(input_shape[::-1])
else:
ret = transpose_fn(a, perm, name=name)
return ret
def batch_scatter(indices, updates, shape, name=None):
with ops.name_scope(name):
indices = ops.convert_to_tensor(indices, name="indices")
indices_shape = array_ops.shape(indices)
indices_dimensions = indices.get_shape().ndims
if indices_dimensions is None:
raise ValueError(
"batch_gather does not allow indices with unknown "
"shape.")
nd_indices = array_ops.expand_dims(indices, axis=-1)
nd_indices_list = []
# Scatter ND requires indices to have an additional dimension, in which the
# coordinates of the updated things are specified. For this to be adapted to
# the scatter_update with several leading dimensions, we simply make use of
# a tf.range for all the leading dimensions followed by concat of all the
# coordinates we created with the original indices.
# For example if indices.shape = [2, 3], we should generate the following
# indices for tf.scatter_nd_update:
# nd_indices[:, :, 0] = [[0, 0, 0], [1, 1, 1]]
# nd_indices[:, :, 1] = [[0, 1, 2], [0, 1, 2]]
# nd_indices[:, :, 2] = indices
for dimension in range(indices_dimensions - 1):
# In this loop we generate the following for the example (one for each
# iteration).
# nd_indices[:, :, 0] = [[0, 0, 0], [1, 1, 1]]
# nd_indices[:, :, 1] = [[0, 1, 2], [0, 1, 2]]
# This is done at every iteration with a tf.range over the size of the
# i-th dimension and using broadcasting over the desired shape.
dimension_size = indices_shape[dimension]
shape_to_broadcast = [1] * (indices_dimensions + 1)
shape_to_broadcast[dimension] = dimension_size
dimension_range = array_ops.reshape(
gen_math_ops._range(0, dimension_size, 1), shape_to_broadcast)
if dimension_range.dtype.base_dtype != nd_indices.dtype:
dimension_range = gen_math_ops.cast(dimension_range,
nd_indices.dtype)
nd_indices_list.append(dimension_range *
array_ops.ones_like(nd_indices))
# Add the original indices at the end, as described above, and concat.
nd_indices_list.append(nd_indices)
final_indices = array_ops.concat(nd_indices_list, axis=-1)
return tf.scatter_nd(final_indices, updates, shape)
def range(start, limit=None, delta=1, name="range"):
"""Creates a sequence of integers.
Creates a sequence of integers that begins at `start` and extends by
increments of `delta` up to but not including `limit`.
Like the Python builtin `range`, `start` defaults to 0, so that
`range(n) = range(0, n)`.
For example:
```
# 'start' is 3
# 'limit' is 18
# 'delta' is 3
tf.range(start, limit, delta) ==> [3, 6, 9, 12, 15]
# 'limit' is 5
tf.range(limit) ==> [0, 1, 2, 3, 4]
```
Args:
start: A 0-D (scalar) of type `int32`. First entry in sequence.
Defaults to 0.
limit: A 0-D (scalar) of type `int32`. Upper limit of sequence,
exclusive.
delta: A 0-D `Tensor` (scalar) of type `int32`. Optional. Default is 1.
Number that increments `start`.
name: A name for the operation (optional).
Returns:
An 1-D `int32` `Tensor`.
"""
if limit is None:
start, limit = 0, start
return gen_math_ops._range(start, limit, delta, name=name)
def batch_scatter_update(ref, indices, updates, use_locking=True, name=None):
"""Generalization of `tf.scatter_update` to axis different than 0.
Analogous to `batch_gather`. This assumes that `ref`, `indices` and `updates`
have a series of leading dimensions that are the same for all of them, and the
updates are performed on the last dimension of indices. In other words, the
dimensions should be the following:
`num_prefix_dims = indices.ndims - 1`
`batch_dim = num_prefix_dims + 1`
`updates.shape = indices.shape + var.shape[batch_dim:]`
where
`updates.shape[:num_prefix_dims]`
`== indices.shape[:num_prefix_dims]`
`== var.shape[:num_prefix_dims]`
And the operation performed can be expressed as:
`var[i_1, ..., i_n, indices[i_1, ..., i_n, j]] = updates[i_1, ..., i_n, j]`
When indices is a 1D tensor, this operation is equivalent to
`tf.scatter_update`.
To avoid this operation there would be 2 alternatives:
1) Reshaping the variable by merging the first `ndims` dimensions. However,
this is not possible because `tf.reshape` returns a Tensor, which we
cannot use `tf.scatter_update` on.
2) Looping over the first `ndims` of the variable and using
`tf.scatter_update` on the subtensors that result of slicing the first
dimension. This is a valid option for `ndims = 1`, but less efficient than
this implementation.
See also `tf.scatter_update` and `tf.scatter_nd_update`.
Args:
ref: `Variable` to scatter onto.
indices: Tensor containing indices as described above.
updates: Tensor of updates to apply to `ref`.
use_locking: Boolean indicating whether to lock the writing operation.
name: Optional scope name string.
Returns:
Ref to `variable` after it has been modified.
Raises:
ValueError: If the initial `ndims` of `ref`, `indices`, and `updates` are
not the same.
"""
with ops.name_scope(name):
indices = ops.convert_to_tensor(indices, name="indices")
indices_shape = array_ops.shape(indices)
indices_dimensions = indices.get_shape().ndims
if indices_dimensions is None:
raise ValueError(
"batch_gather does not allow indices with unknown "
"shape.")
nd_indices = array_ops.expand_dims(indices, axis=-1)
nd_indices_list = []
# Scatter ND requires indices to have an additional dimension, in which the
# coordinates of the updated things are specified. For this to be adapted to
# the scatter_update with several leading dimensions, we simply make use of
# a tf.range for all the leading dimensions followed by concat of all the
# coordinates we created with the original indices.
# For example if indices.shape = [2, 3, 4], we should generate the following
# indices for tf.scatter_nd_update:
# nd_indices[:, :, 0] = [[0, 0, 0], [1, 1, 1]]
# nd_indices[:, :, 1] = [[0, 1, 2], [0, 1, 2]]
# nd_indices[:, :, 2] = indices
for dimension in range(indices_dimensions - 1):
# In this loop we generate the following for the example (one for each
# iteration).
# nd_indices[:, :, 0] = [[0, 0, 0], [1, 1, 1]]
# nd_indices[:, :, 1] = [[0, 1, 2], [0, 1, 2]]
# This is done at every iteration with a tf.range over the size of the
# i-th dimension and using broadcasting over the desired shape.
dimension_size = indices_shape[dimension]
shape_to_broadcast = [1] * (indices_dimensions + 1)
shape_to_broadcast[dimension] = dimension_size
dimension_range = array_ops.reshape(
gen_math_ops._range(0, dimension_size, 1), shape_to_broadcast)
if dimension_range.dtype.base_dtype != nd_indices.dtype:
dimension_range = gen_math_ops.cast(dimension_range,
nd_indices.dtype)
nd_indices_list.append(dimension_range *
array_ops.ones_like(nd_indices))
# Add the original indices at the end, as described above, and concat.
nd_indices_list.append(nd_indices)
final_indices = array_ops.concat(nd_indices_list, axis=-1)
return scatter_nd_update(ref,
final_indices,
updates,
use_locking=use_locking)
def batch_gather(params, indices, axis, name=None):
"""
Extension of the batch_gather function in tensorflow
(see https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/array_ops.py
or https://www.tensorflow.org/api_docs/python/tf/batch_gather)
Gather slices from `params` according to `indices` with leading batch dims.
This operation assumes that the leading dimensions of `indices` are dense,
and the gathers on the axis corresponding to the last dimension of `indices`.
More concretely it computes:
`result[i1, ..., in, j1, ..., jm, k1, ...., kl] = params[i1, ..., in, indices[i1, ..., in, j1, ..., jm], k1, ..., kl]`
Therefore `params` should be a Tensor of shape [A1, ..., AN, C0, B1, ..., BM],
`indices` should be a Tensor of shape [A1, ..., AN, C1, ..., CK] and `result` will be
a Tensor of size `[A1, ..., AN, C1, ..., CK, B1, ..., BM]`.
In the case in which indices is a 1D tensor, this operation is equivalent to
`tf.gather`.
See also `tf.gather` and `tf.gather_nd`.
Args:
params: A `Tensor`. The tensor from which to gather values.
indices: A `Tensor`. Must be one of the following types: int32, int64. Index
tensor. Must be in range `[0, params.shape[axis]`, where `axis` is the
last dimension of `indices` itself.
axis: A `Tensor`. Must be one of the following types: int32, int64. The axis
in `params` to gather `indices` from.
name: A name for the operation (optional).
Returns:
A Tensor. Has the same type as `params`.
Raises:
ValueError: if `indices` has an unknown shape.
"""
with ops.name_scope(name):
indices = ops.convert_to_tensor(indices, name="indices")
params = ops.convert_to_tensor(params, name="params")
indices_shape = tf.shape(indices)
params_shape = tf.shape(params)
ndims = indices.shape.ndims
if ndims is None:
raise ValueError(
"batch_gather does not allow indices with unknown "
"shape.")
batch_indices = indices
indices_dtype = indices.dtype.base_dtype
accum_dim_value = tf.ones((), dtype=indices_dtype)
# Use correct type for offset index computation
casted_params_shape = gen_math_ops.cast(params_shape, indices_dtype)
for dim in range(axis, 0, -1):
dim_value = casted_params_shape[dim - 1]
accum_dim_value *= casted_params_shape[dim]
start = tf.zeros((), dtype=indices_dtype)
step = tf.ones((), dtype=indices_dtype)
dim_indices = gen_math_ops._range(start, dim_value, step)
dim_indices *= accum_dim_value
dim_shape = tf.stack([1] * (dim - 1) + [dim_value] + [1] *
(ndims - dim),
axis=0)
batch_indices += tf.reshape(dim_indices, dim_shape)
flat_inner_shape_indices = gen_math_ops.prod(
indices_shape[:(axis + 1)], [0], False)
flat_indices = tf.reshape(
batch_indices,
tf.concat([[flat_inner_shape_indices], indices_shape[(axis + 1):]],
axis=0))
outer_shape = params_shape[(axis + 1):]
flat_inner_shape_params = gen_math_ops.prod(params_shape[:(axis + 1)],
[0], False)
flat_params = tf.reshape(
params, tf.concat([[flat_inner_shape_params], outer_shape],
axis=0))
flat_result = tf.gather(flat_params, flat_indices)
result = tf.reshape(flat_result,
tf.concat([indices_shape, outer_shape], axis=0))
final_shape = indices.get_shape()[:axis].merge_with(
params.get_shape()[:axis])
final_shape = final_shape.concatenate(indices.get_shape()[axis:])
final_shape = final_shape.concatenate(params.get_shape()[(axis + 1):])
result.set_shape(final_shape)
return result
def batch_scatter_update(ref, indices, updates, use_locking=True, name=None):
"""Generalization of `tf.scatter_update` to axis different than 0.
Analogous to `batch_gather`. This assumes that `ref`, `indices` and `updates`
have a series of leading dimensions that are the same for all of them, and the
updates are performed on the last dimension of indices. In other words, the
dimensions should be the following:
`num_prefix_dims = indices.ndims - 1`
`batch_dim = num_prefix_dims + 1`
`updates.shape = indices.shape + var.shape[batch_dim:]`
where
`updates.shape[:num_prefix_dims]`
`== indices.shape[:num_prefix_dims]`
`== var.shape[:num_prefix_dims]`
And the operation performed can be expressed as:
`var[i_1, ..., i_n, indices[i_1, ..., i_n, j]] = updates[i_1, ..., i_n, j]`
When indices is a 1D tensor, this operation is equivalent to
`tf.scatter_update`.
To avoid this operation there would be 2 alternatives:
1) Reshaping the variable by merging the first `ndims` dimensions. However,
this is not possible because `tf.reshape` returns a Tensor, which we
cannot use `tf.scatter_update` on.
2) Looping over the first `ndims` of the variable and using
`tf.scatter_update` on the subtensors that result of slicing the first
dimension. This is a valid option for `ndims = 1`, but less efficient than
this implementation.
See also `tf.scatter_update` and `tf.scatter_nd_update`.
Args:
ref: `Variable` to scatter onto.
indices: Tensor containing indices as described above.
updates: Tensor of updates to apply to `ref`.
use_locking: Boolean indicating whether to lock the writing operation.
name: Optional scope name string.
Returns:
Ref to `variable` after it has been modified.
Raises:
ValueError: If the initial `ndims` of `ref`, `indices`, and `updates` are
not the same.
"""
with ops.name_scope(name):
indices = ops.convert_to_tensor(indices, name="indices")
indices_shape = array_ops.shape(indices)
indices_dimensions = indices.get_shape().ndims
if indices_dimensions is None:
raise ValueError("batch_gather does not allow indices with unknown "
"shape.")
nd_indices = array_ops.expand_dims(indices, axis=-1)
nd_indices_list = []
# Scatter ND requires indices to have an additional dimension, in which the
# coordinates of the updated things are specified. For this to be adapted to
# the scatter_update with several leading dimensions, we simply make use of
# a tf.range for all the leading dimensions followed by concat of all the
# coordinates we created with the original indices.
# For example if indices.shape = [2, 3, 4], we should generate the following
# indices for tf.scatter_nd_update:
# nd_indices[:, :, 0] = [[0, 0, 0], [1, 1, 1]]
# nd_indices[:, :, 1] = [[0, 1, 2], [0, 1, 2]]
# nd_indices[:, :, 2] = indices
for dimension in range(indices_dimensions - 1):
# In this loop we generate the following for the example (one for each
# iteration).
# nd_indices[:, :, 0] = [[0, 0, 0], [1, 1, 1]]
# nd_indices[:, :, 1] = [[0, 1, 2], [0, 1, 2]]
# This is done at every iteration with a tf.range over the size of the
# i-th dimension and using broadcasting over the desired shape.
dimension_size = indices_shape[dimension]
shape_to_broadcast = [1] * (indices_dimensions + 1)
shape_to_broadcast[dimension] = dimension_size
dimension_range = array_ops.reshape(
gen_math_ops._range(0, dimension_size, 1), shape_to_broadcast)
if dimension_range.dtype.base_dtype != nd_indices.dtype:
dimension_range = gen_math_ops.cast(dimension_range, nd_indices.dtype)
nd_indices_list.append(
dimension_range * array_ops.ones_like(nd_indices))
# Add the original indices at the end, as described above, and concat.
nd_indices_list.append(nd_indices)
final_indices = array_ops.concat(nd_indices_list, axis=-1)
return scatter_nd_update(
ref, final_indices, updates, use_locking=use_locking)
def call(self, inputs, state):
"""VCGRU with nunits cells."""
add = math_ops.add
sub = math_ops.subtract
mult = math_ops.multiply
# computing m_t
m_t = add(math_ops.matmul(state, self._u),
math_ops.matmul(inputs, self._v))
m_t = nn_ops.bias_add(m_t, self._b)
m_t = math_ops.sigmoid(m_t)
# add L1 loss
ops.add_to_collection('L1 loss', math_ops.abs(m_t - self._m_target))
# computing e_t (= thr)
i = gen_math_ops._range(1, self._num_units + 1, 1)
i = math_ops.cast(i, dtype=dtypes.float32)
mtD = gen_array_ops.tile(mult(m_t[1], self._num_units),
[self._num_units])
thr = math_ops.sigmoid(mult(self._sharpness, sub(mtD, i)))
thr = math_ops.round(add(thr, sub(0.5, self._epsilon)))
ones = array_ops.ones_like(thr)
thr_inv = sub(ones, thr)
# computing h_t
if inputs.shape[1] < thr.shape[0]:
_inputs = mult(inputs,
array_ops.slice(thr, [
0,
], [
inputs.shape[1],
]))
elif inputs.shape[1] > thr.shape[0]:
_inputs = mult(
inputs,
array_ops.concat(1, [
thr,
array_ops.zeros_like(inputs.shape[1] - thr.shape[0])
]))
else:
_inputs = mult(inputs, thr)
_state = mult(state, thr)
gate_inputs = math_ops.matmul(array_ops.concat([_inputs, _state], 1),
self._gate_kernel)
gate_inputs = nn_ops.bias_add(gate_inputs, self._gate_bias)
value = math_ops.sigmoid(gate_inputs)
r, u = array_ops.split(value=value, num_or_size_splits=2, axis=1)
_r_state = r * _state
candidate = math_ops.matmul(array_ops.concat([_inputs, _r_state], 1),
self._candidate_kernel)
candidate = nn_ops.bias_add(candidate, self._candidate_bias)
c = self._activation(candidate)
new_h = u * _state + (1 - u) * c
output = add(mult(new_h, thr), mult(state, thr_inv))
return output, output
