def _BatchGatherGrad(params_shape, values, indices, batch_dims,
gather_dim_size):
"""Returns the gradient of GatherV2 with batch dimensions."""
# Axis is the first non-batch dimension.
indices_size = array_ops.expand_dims(array_ops.size(indices), 0)
if batch_dims:
values_shape = array_ops.shape(values)
# Add the batch offsets to indices and flatten the batch dimensions.
outer_shape = values_shape[:batch_dims]
inner_shape = values_shape[batch_dims:][1:]
batch_size = gen_math_ops.prod(outer_shape, [0], False)
flat_values_shape = array_ops.concat([[-1], inner_shape], 0)
gather_dim_size *= batch_size
indices = _GetBatchIndices(params_shape, indices, batch_dims)
values = array_ops.reshape(_IndexedSlicesToTensorNoWarning(values),
flat_values_shape)
indices = array_ops.reshape(indices, indices_size)
params_grad = math_ops.unsorted_segment_sum(values, indices,
gather_dim_size)
if batch_dims:
# Put back the batch dimensions.
params_grad = array_ops.reshape(
params_grad, array_ops.concat([outer_shape, flat_values_shape], 0))
return params_grad
def size_internal(input, name=None, optimize=True, out_type=dtypes.int32):
if context.executing_eagerly() and not isinstance(
input,
(sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)):
input = ops.convert_to_tensor(input)
np_out_type = out_type.as_numpy_dtype
num_elements = np.prod(input._shape_tuple(), dtype=np_out_type)
return ops.convert_to_tensor(num_elements, dtype=out_type)
with ops.name_scope(name, "Size", [input]) as name:
if isinstance(
input,
(sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)):
return gen_math_ops.prod(gen_math_ops.cast(input.dense_shape,
out_type),
0,
name=name)
else:
input_tensor = ops.convert_to_tensor(input)
input_shape = input_tensor.get_shape()
if optimize:
if input_shape.is_fully_defined():
return constant(input_shape.num_elements(),
out_type,
name=name)
if input_shape.dims and any(dim == 0
for dim in input_shape.dims):
return constant(0, out_type, name=name)
return gen_array_ops.size(input, name=name, out_type=out_type)
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 tf.name_scope(name):
indices = tf.convert_to_tensor(indices, name="indices")
params = tf.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 = tf.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.reshape(dim_indices, dim_shape)
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 call(self, inputs):
def _bucketize_op(bins):
bins = [math_ops.cast(bins, dtypes.float32)]
return lambda inputs: gen_boosted_trees_ops.BoostedTreesBucketize( # pylint: disable=g-long-lambda
float_values=[math_ops.cast(inputs, dtypes.float32)],
bucket_boundaries=bins)[0]
if tf_utils.is_ragged(inputs):
integer_buckets = ragged_functional_ops.map_flat_values(
_bucketize_op(array_ops.squeeze(self.bins)), inputs)
# Ragged map_flat_values doesn't touch the non-values tensors in the
# ragged composite tensor. If this op is the only op a Keras model,
# this can cause errors in Graph mode, so wrap the tensor in an identity.
return array_ops.identity(integer_buckets)
elif isinstance(inputs, sparse_tensor.SparseTensor):
integer_buckets = gen_boosted_trees_ops.BoostedTreesBucketize(
float_values=[math_ops.cast(inputs.values, dtypes.float32)],
bucket_boundaries=[
math_ops.cast(array_ops.squeeze(self.bins), dtypes.float32)
])[0]
return sparse_tensor.SparseTensor(
indices=array_ops.identity(inputs.indices),
values=integer_buckets,
dense_shape=array_ops.identity(inputs.dense_shape))
else:
input_shape = inputs.get_shape()
if any(dim is None for dim in input_shape.as_list()[1:]):
raise NotImplementedError(
"Discretization Layer requires known non-batch shape,"
"found {}".format(input_shape))
reshaped = array_ops.reshape(
inputs,
[-1, gen_math_ops.prod(input_shape.as_list()[1:], axis=0)])
return array_ops.reshape(
control_flow_ops.vectorized_map(
_bucketize_op(array_ops.squeeze(self.bins)), reshaped),
array_ops.constant([-1] + input_shape.as_list()[1:]))
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
