def testError(self,
descr,
mode,
data,
repeats,
axis,
exception=ValueError,
error=None):
# Make sure that this is also an error case for numpy.
with self.assertRaises(exception):
np.repeat(data, repeats, axis)
if mode == 'constant':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
elif mode == 'dynamic':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
data = array_ops.placeholder_with_default(data, data.shape)
repeats = array_ops.placeholder_with_default(repeats, repeats.shape)
elif mode == 'unknown_shape':
data = array_ops.placeholder_with_default(data, None)
repeats = array_ops.placeholder_with_default(repeats, None)
with self.assertRaisesRegexp(exception, error):
ragged_util.repeat(data, repeats, axis)
def testError(self,
descr,
mode,
data,
repeats,
axis,
exception=ValueError,
error=None):
# Make sure that this is also an error case for numpy.
with self.assertRaises(exception):
np.repeat(data, repeats, axis)
if mode == 'constant':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
elif mode == 'dynamic':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
data = array_ops.placeholder_with_default(data, data.shape)
repeats = array_ops.placeholder_with_default(
repeats, repeats.shape)
elif mode == 'unknown_shape':
data = array_ops.placeholder_with_default(data, None)
repeats = array_ops.placeholder_with_default(repeats, None)
with self.assertRaisesRegex(exception, error):
ragged_util.repeat(data, repeats, axis)
def combine_graphs(batch):
"""Combine multiple graphs into a single network"""
# Compute the mappings from bond index to graph index
batch_size = tf.size(batch['n_atom'], name='batch_size')
mol_id = tf.range(batch_size, name='mol_inds')
batch['node_graph_indices'] = repeat(mol_id, batch['n_atom'], axis=0)
batch['bond_graph_indices'] = repeat(mol_id, batch['n_bond'], axis=0)
# Reshape the bond, connectivity, and node lists
for c in ['atom', 'bond', 'connectivity']:
batch[c] = batch[c].flat_values
# Reshape the connectivity matrix to (None, 2)
batch['connectivity'] = tf.reshape(batch['connectivity'], (-1, 2))
# Denote the shapes for the atom and bond matrices
# Only an issue for 1.14, which cannot infer them it seems
for c in ['atom', 'bond']:
batch[c].set_shape((None,))
# Compute offsets for the connectivity matrix
offset_values = tf.cumsum(batch['n_atom'], exclusive=True)
offsets = repeat(offset_values, batch['n_bond'], name='offsets', axis=0)
batch['connectivity'] += tf.expand_dims(offsets, 1)
return batch
def repeat(values, repeats, axis):
"""See https://github.com/tensorflow/tensorflow/issues/8246."""
try:
repeat = tf.repeat
except AttributeError:
from tensorflow.python.ops.ragged.ragged_util import repeat # pylint: disable=import-error
if values.shape.ndims == 1:
return repeat(values, repeats, axis=axis)
else:
indices = repeat(tf.range(tf.shape(values)[axis]), repeats, axis=0)
return tf.gather(values, indices, axis=axis)
def row_splits_to_segment_ids(splits, name=None):
"""Generates the segmentation corresponding to a RaggedTensor `splits` vector.
Returns an integer vector `segment_ids`, where `segment_ids[i] == j` if
`splits[j] <= i < splits[j+1]`. Example:
```python
>>> ragged.row_splits_to_segment_ids([0, 3, 3, 5, 6, 9]).eval()
[ 0 0 0 2 2 3 4 4 4 ]
```
Args:
splits: A sorted 1-D int64 Tensor. `splits[0]` must be zero.
name: A name prefix for the returned tensor (optional).
Returns:
A sorted 1-D int64 Tensor, with `shape=[splits[-1]]`
Raises:
ValueError: If `splits` is invalid.
"""
with ops.name_scope(name, "RaggedSplitsToSegmentIds", [splits]) as name:
splits = ops.convert_to_tensor(splits,
dtype=dtypes.int64,
name="splits")
splits.shape.assert_has_rank(1)
if tensor_shape.dimension_value(splits.shape[0]) == 0:
raise ValueError("Invalid row_splits: []")
row_lengths = splits[1:] - splits[:-1]
nrows = array_ops.shape(splits, out_type=dtypes.int64)[-1] - 1
indices = math_ops.range(nrows)
return ragged_util.repeat(indices, repeats=row_lengths, axis=0)
def row_splits_to_segment_ids(splits, name=None):
"""Generates the segmentation corresponding to a RaggedTensor `row_splits`.
Returns an integer vector `segment_ids`, where `segment_ids[i] == j` if
`splits[j] <= i < splits[j+1]`. Example:
```python
>>> ragged.row_splits_to_segment_ids([0, 3, 3, 5, 6, 9]).eval()
[ 0 0 0 2 2 3 4 4 4 ]
```
Args:
splits: A sorted 1-D int64 Tensor. `splits[0]` must be zero.
name: A name prefix for the returned tensor (optional).
Returns:
A sorted 1-D int64 Tensor, with `shape=[splits[-1]]`
Raises:
ValueError: If `splits` is invalid.
"""
with ops.name_scope(name, "RaggedSplitsToSegmentIds", [splits]) as name:
splits = ops.convert_to_tensor(splits, dtype=dtypes.int64, name="splits")
splits.shape.assert_has_rank(1)
if tensor_shape.dimension_value(splits.shape[0]) == 0:
raise ValueError("Invalid row_splits: []")
row_lengths = splits[1:] - splits[:-1]
nrows = array_ops.shape(splits, out_type=dtypes.int64)[-1] - 1
indices = math_ops.range(nrows)
return ragged_util.repeat(indices, repeats=row_lengths, axis=0)
def _expand_and_tile(global_features, row_splits_or_k):
if not isinstance(row_splits_or_k, tf.Tensor) or (
row_splits_or_k.shape.ndims == 0):
# knn
raise NotImplementedError('TODO')
else:
from tensorflow.python.ops.ragged.ragged_util import repeat
return repeat(global_features, row_splits_or_k, axis=0)
def tf_repeat(values, repeats):
values = tf.convert_to_tensor(values)
repeats = tf.convert_to_tensor(repeats)
if values.shape.ndims != 1:
raise ValueError("values must be rank 1, got shape %s" % values.shape)
if repeats.shape.ndims != 1:
raise ValueError("repeats must be rank 1, got shape %s" % repeats.shape)
if not repeats.dtype.is_integer:
raise ValueError("repeats must be an integer, got %s" % repeats.dtype)
try:
from tensorflow.python.ops.ragged.ragged_util import repeat
return repeat(values, repeats, axis=0)
except ImportError:
return _foldl_repeat(values, repeats)
def testValuesMatchesNumpy(self, mode, data, repeats, axis):
# Exception: we can't handle negative axis if data.ndims is unknown.
if axis < 0 and mode == 'unknown_shape':
return
expected = np.repeat(data, repeats, axis)
if mode == 'constant':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
elif mode == 'dynamic':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
data = array_ops.placeholder_with_default(data, data.shape)
repeats = array_ops.placeholder_with_default(repeats, repeats.shape)
elif mode == 'unknown_shape':
data = array_ops.placeholder_with_default(data, None)
repeats = array_ops.placeholder_with_default(repeats, None)
result = ragged_util.repeat(data, repeats, axis)
self.assertAllEqual(result, expected)
def testValuesMatchesNumpy(self, mode, data, repeats, axis):
# Exception: we can't handle negative axis if data.ndims is unknown.
if axis < 0 and mode == 'unknown_shape':
return
expected = np.repeat(data, repeats, axis)
if mode == 'constant':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
elif mode == 'dynamic':
data = constant_op.constant(data)
repeats = constant_op.constant(repeats)
data = array_ops.placeholder_with_default(data, data.shape)
repeats = array_ops.placeholder_with_default(repeats, repeats.shape)
elif mode == 'unknown_shape':
data = array_ops.placeholder_with_default(data, None)
repeats = array_ops.placeholder_with_default(repeats, None)
result = ragged_util.repeat(data, repeats, axis)
self.assertAllEqual(result, expected)
def flat_expanding_global_deconv(
global_features, coord_features, row_splits_or_k):
"""
Global deconvolution operation.
Args:
global_features: [pi, fi]
coord_features: [po, fk]
row_splits_or_k: [pi+1]
Returns:
convolved features: [po, fi*fk]
"""
from tensorflow.python.ops.ragged.ragged_util import repeat
if row_splits_or_k.shape.ndims == 0:
raise NotImplementedError
global_features = repeat(
global_features, utils.diff(row_splits_or_k), axis=0)
merged = utils.merge_features(global_features, coord_features)
merged = utils.flatten_final_dims(merged, 2)
return merged
def row_splits_to_segment_ids(splits, name=None, out_type=None):
"""Generates the segmentation corresponding to a RaggedTensor `row_splits`.
Returns an integer vector `segment_ids`, where `segment_ids[i] == j` if
`splits[j] <= i < splits[j+1]`. Example:
>>> print(tf.ragged.row_splits_to_segment_ids([0, 3, 3, 5, 6, 9]))
tf.Tensor([0 0 0 2 2 3 4 4 4], shape=(9,), dtype=int64)
Args:
splits: A sorted 1-D integer Tensor. `splits[0]` must be zero.
name: A name prefix for the returned tensor (optional).
out_type: The dtype for the return value. Defaults to `splits.dtype`,
or `tf.int64` if `splits` does not have a dtype.
Returns:
A sorted 1-D integer Tensor, with `shape=[splits[-1]]`
Raises:
ValueError: If `splits` is invalid.
"""
with ops.name_scope(name, "RaggedSplitsToSegmentIds", [splits]) as name:
splits = ops.convert_to_tensor(splits,
name="splits",
preferred_dtype=dtypes.int64)
if splits.dtype not in (dtypes.int32, dtypes.int64):
raise ValueError("splits must have dtype int32 or int64")
splits.shape.assert_has_rank(1)
if tensor_shape.dimension_value(splits.shape[0]) == 0:
raise ValueError("Invalid row_splits: []")
if out_type is None:
out_type = splits.dtype
else:
out_type = dtypes.as_dtype(out_type)
row_lengths = splits[1:] - splits[:-1]
nrows = array_ops.shape(splits, out_type=out_type)[-1] - 1
indices = math_ops.range(nrows)
return ragged_util.repeat(indices, repeats=row_lengths, axis=0)
def testRepeat(self, data, repeats, expected, axis=None):
result = ragged_util.repeat(data, repeats, axis)
self.assertAllEqual(result, expected)
def testRepeat(self, data, repeats, expected, axis=None): result = ragged_util.repeat(data, repeats, axis) self.assertAllEqual(result, expected)
def _class_maps_to_colored_maps(self, class_map):
gray_gts = tf.argmax(class_map, axis=3,
output_type=tf.int32) * (255 // self.n_classes)
temp = tf.expand_dims(gray_gts, 3)
rgb_gt = repeat(temp, 3, axis=3)
return tf.cast(rgb_gt, dtype=tf.uint8)
def testRepeat(self, data, repeats, expected, axis=None):
result = ragged_util.repeat(data, repeats, axis)
with self.test_session():
self.assertEqual(result.eval().tolist(), expected)
def _repeat(args, axis=0):
from tensorflow.python.ops.ragged.ragged_util import repeat
data, repeats = args
return repeat(data, repeats, axis=axis)
def batch_gather(params, indices, name=None):
"""Gathers slices from `params` according to `indices` with batch dims.
This operation is similar to `gather`, but it assumes that the leading `N`
dimensions of `indices` and `params` are batch dimensions, and performs a
gather within each batch. In particular, when using this operation with `N`
batch dimensions `B1...BN`:
* `indices` has shape `[B1...BN, I]`
* `params` has shape `[B1...BN, P1...PM]`.
* `result` has shape `[B1...BN, I, P2...PM]`.
* `result[b1...bN, i, p2...pM] =
params[b1...bN, indices[b1...bN, i], p2...pM]`
Args:
params: A potentially ragged tensor with shape `[B1...BN, P1...PM]` (`N>=0`,
`M>0`).
indices: A potentially ragged tensor with shape `[B1...BN, I]` (`N>=0`).
name: A name for the operation (optional).
Returns:
A potentially ragged tensor with shape `[B1...BN, I, P2...PM]`.
`result.ragged_rank = max(indices.ragged_rank, params.ragged_rank)`.
#### Example:
```python
>>> params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
>>> indices = tf.ragged.constant([[1, 2, 0], [], [], [0, 0]])
>>> tf.compat.v1.batch_gather(params, indices)
[['b', 'c', 'a'], [], [], ['e', 'e']]
```
"""
if not (ragged_tensor.is_ragged(params) or ragged_tensor.is_ragged(indices)):
return array_ops.batch_gather(params, indices, name)
with ops.name_scope(name, 'RaggedBatchGather', [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_ndims = indices.shape.ndims
if indices_ndims is None:
raise ValueError(
'batch_gather does not allow indices with unknown shape.')
if indices_ndims == 0:
raise ValueError('indices.rank must be at least 1.')
if ragged_tensor.is_ragged(indices):
# If the outermost ragged dimension is a batch dimension, recurse.
if indices_ndims > 2:
if not ragged_tensor.is_ragged(params):
raise ValueError('batch shape from indices does '
'not match params shape')
checks = [check_ops.assert_equal(params.row_splits, indices.row_splits)]
with ops.control_dependencies(checks):
return ragged_tensor.RaggedTensor.from_row_splits(
batch_gather(params.values, indices.values), indices.row_splits,
validate=False)
# Otherwise, indices is a 2D ragged tensor with 1 ragged dimension.
else:
# Ensure that `params` is ragged and has at least 2 dimensions.
if not ragged_tensor.is_ragged(params):
if params.shape.ndims is not None and params.shape.ndims < 2:
raise ValueError('batch shape from indices does '
'not match params shape')
params = ragged_tensor.RaggedTensor.from_tensor(
params, ragged_rank=1,
row_splits_dtype=indices.row_splits.dtype)
# Adjust indices from within-batch to global (in params.values), and
# then use ragged.gather to gather them.
num_indices = indices.row_lengths()
params_starts = params.row_starts()
adjustments = ragged_util.repeat(params_starts, num_indices, axis=0)
adjusted_index_values = (
math_ops.cast(indices.values, adjustments.dtype) + adjustments)
return ragged_tensor.RaggedTensor.from_row_splits(
ragged_gather_ops.gather(params.values, adjusted_index_values),
indices.row_splits, validate=False)
else: # params is a RaggedTensor and indices is a Tensor.
if indices_ndims == 1:
return ragged_gather_ops.gather(params, indices)
elif indices_ndims == 2:
# Adjust indices from batch-local to global (in params.values)
adjustments = array_ops.expand_dims(params.row_starts(), 1)
adjusted_indices = (
math_ops.cast(indices, adjustments.dtype) + adjustments)
return ragged_gather_ops.gather(params.values, adjusted_indices)
else:
raise ValueError('batch shape from indices does not match params shape')
def batch_gather(params, indices, name=None):
"""Gathers slices from `params` according to `indices` with batch dims.
This operation is similar to `gather`, but it assumes that the leading `N`
dimensions of `indices` and `params` are batch dimensions, and performs a
gather within each batch. In particular, when using this operation with `N`
batch dimensions `B1...BN`:
* `indices` has shape `[B1...BN, I]`
* `params` has shape `[B1...BN, P1...PM]`.
* `result` has shape `[B1...BN, I, P2...PM]`.
* `result[b1...bN, i, p2...pM] =
params[b1...bN, indices[b1...bN, i], p2...pM]`
Args:
params: A potentially ragged tensor with shape `[B1...BN, P1...PM]` (`N>=0`,
`M>0`).
indices: A potentially ragged tensor with shape `[B1...BN, I]` (`N>=0`).
name: A name for the operation (optional).
Returns:
A potentially ragged tensor with shape `[B1...BN, I, P2...PM]`.
`result.ragged_rank = max(indices.ragged_rank, params.ragged_rank)`.
#### Example:
```python
>>> params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
>>> indices = tf.ragged.constant([[1, 2, 0], [], [], [0, 0]])
>>> ragged.batch_gather(params, indices)
[['b', 'c', 'a'], [], [], ['e', 'e']]
```
"""
if not (ragged_tensor.is_ragged(params) or ragged_tensor.is_ragged(indices)):
return array_ops.batch_gather(params, indices, name)
with ops.name_scope(name, 'RaggedBatchGather', [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')
indices_ndims = indices.shape.ndims
if indices_ndims is None:
raise ValueError(
'batch_gather does not allow indices with unknown shape.')
if indices_ndims == 0:
raise ValueError('indices.rank must be at least 1.')
if ragged_tensor.is_ragged(indices):
# If the outermost ragged dimension is a batch dimension, recurse.
if indices_ndims > 2:
if not ragged_tensor.is_ragged(params):
raise ValueError('batch shape from indices does '
'not match params shape')
checks = [check_ops.assert_equal(params.row_splits, indices.row_splits)]
with ops.control_dependencies(checks):
return ragged_tensor.RaggedTensor.from_row_splits(
batch_gather(params.values, indices.values), indices.row_splits)
# Otherwise, indices is a 2D ragged tensor with 1 ragged dimension.
else:
# Ensure that `params` is ragged and has at least 2 dimensions.
if not ragged_tensor.is_ragged(params):
if params.shape.ndims is not None and params.shape.ndims < 2:
raise ValueError('batch shape from indices does '
'not match params shape')
params = ragged_conversion_ops.from_tensor(params, ragged_rank=1)
# Adjust indices from within-batch to global (in params.values), and
# then use ragged.gather to gather them.
num_indices = indices.row_lengths()
params_starts = params.row_starts()
adjustments = ragged_util.repeat(params_starts, num_indices, axis=0)
adjusted_index_values = math_ops.to_int64(indices.values) + adjustments
return ragged_tensor.RaggedTensor.from_row_splits(
gather(params.values, adjusted_index_values), indices.row_splits)
else: # params is a RaggedTensor and indices is a Tensor.
if indices_ndims == 1:
return gather(params, indices)
elif indices_ndims == 2:
# Adjust indices from batch-local to global (in params.values)
adjustments = array_ops.expand_dims(params.row_starts(), 1)
adjusted_indices = math_ops.to_int64(indices) + adjustments
return gather(params.values, adjusted_indices)
else:
raise ValueError('batch shape from indices does not match params shape')
def segmentation_logits(inputs,
output_spec,
num_obj_classes=16,
r0=0.1,
initial_filters=(16, ),
initial_activation=seg_activation,
filters=(32, 64, 128, 256),
global_units='combined',
query_fn=core.query_pairs,
radii_fn=core.constant_radii,
global_deconv_all=False,
coords_transform=None,
weights_transform=None,
convolver=None):
if convolver is None:
convolver = c.ExpandingConvolver(activation=seg_activation)
if coords_transform is None:
coords_transform = t.polynomial_transformer()
if weights_transform is None:
def weights_transform(*args, **kwargs):
return None
coords = inputs['positions']
normals = inputs.get('normals')
if normals is None:
raise NotImplementedError()
features = b.as_batched_model_input(normals)
for f in initial_filters:
features = tf.ragged.map_flat_values(core_layers.Dense(f), features)
features = tf.ragged.map_flat_values(initial_activation, features)
assert (isinstance(features, tf.RaggedTensor)
and features.ragged_rank == 1)
class_embeddings = _get_class_embeddings(
b.as_batched_model_input(inputs['obj_label']), num_obj_classes,
[initial_filters[-1], filters[0]])
features = core.add_local_global(features, class_embeddings[0])
input_row_splits = features.row_splits
features = utils.flatten_leading_dims(features, 2)
n_res = len(filters)
unscaled_radii2 = radii_fn(n_res)
if isinstance(unscaled_radii2, tf.Tensor):
assert (unscaled_radii2.shape == (n_res, ))
radii2 = utils.lambda_call(tf.math.scalar_mul, r0**2, unscaled_radii2)
radii2 = tf.keras.layers.Lambda(tf.unstack,
arguments=dict(axis=0))(radii2)
for i, radius2 in enumerate(radii2):
tf.compat.v1.summary.scalar('r%d' % i,
tf.sqrt(radius2),
family='radii')
else:
radii2 = unscaled_radii2 * (r0**2)
def maybe_feed(r2):
is_tensor_or_var = isinstance(r2, (tf.Tensor, tf.Variable))
if is_tensor_or_var:
return b.prebatch_feed(tf.keras.layers.Lambda(tf.sqrt)(radius2))
else:
return np.sqrt(r2)
pp_radii2 = [maybe_feed(r2) for r2 in radii2]
all_features = []
in_place_neighborhoods = []
sampled_neighborhoods = []
global_features = []
# encoder
for i, (radius2, pp_radius2) in enumerate(zip(radii2, pp_radii2)):
neighbors, sample_rate = query_fn(coords,
pp_radius2,
name='query%d' % i)
if not isinstance(radius2, tf.Tensor):
radius2 = utils.constant(radius2, dtype=tf.float32)
neighborhood = n.InPlaceNeighborhood(coords, neighbors)
in_place_neighborhoods.append(neighborhood)
features, nested_row_splits = core.convolve(features, radius2,
filters[i], neighborhood,
coords_transform,
weights_transform,
convolver.in_place_conv)
all_features.append(features)
if global_units == 'combined':
coord_features = coords_transform(neighborhood.out_coords, None)
global_features.append(
convolver.global_conv(features, coord_features,
nested_row_splits[-2], filters[i]))
global_features = tf.keras.layers.Lambda(
tf.concat, arguments=dict(axis=-1))(global_features)
# resample
if i < n_res - 1:
sample_indices = sample.sample(
sample_rate,
tf.keras.layers.Lambda(lambda s: tf.size(s) // 4)(sample_rate))
neighborhood = n.SampledNeighborhood(neighborhood, sample_indices)
sampled_neighborhoods.append(neighborhood)
features, nested_row_splits = core.convolve(
features, radius2, filters[i + 1], neighborhood,
coords_transform, weights_transform, convolver.resample_conv)
coords = neighborhood.out_coords
# global_conv
if global_units is not None:
row_splits = nested_row_splits[-2]
if global_units == 'combined':
global_features = tf.keras.layers.Lambda(
tf.concat, arguments=dict(axis=-1))(global_features)
else:
coord_features = coords_transform(coords, None)
global_features = convolver.global_conv(features, coord_features,
row_splits, global_units)
coord_features = coords_transform(coords, None)
features = convolver.global_deconv(global_features, coord_features,
row_splits, filters[-1])
# decoder
for i in range(n_res - 1, -1, -1):
if i < n_res - 1:
# up-sample
neighborhood = sampled_neighborhoods.pop().transpose
features, nested_row_splits = core.convolve(
features, radius2, filters[i], neighborhood, coords_transform,
weights_transform, convolver.resample_conv)
if global_deconv_all:
coords = neighborhood.out_coords
row_splits = \
neighborhood.offset_batched_neighbors.nested_row_splits[-2]
coord_features = coords_transform(coords)
deconv_features = convolver.global_deconv(
global_features, coord_features, row_splits, filters[i])
features = tf.keras.layers.Add()([features, deconv_features])
forward_features = all_features.pop()
if not (i == n_res - 1 and global_units is None):
features = tf.keras.layers.Lambda(tf.concat,
arguments=dict(axis=-1))(
[features, forward_features])
neighborhood = in_place_neighborhoods.pop().transpose
features, nested_row_splits = core.convolve(features, radius2,
filters[i], neighborhood,
coords_transform,
weights_transform,
convolver.resample_conv)
features = tf.RaggedTensor.from_row_splits(features, input_row_splits)
features = core.add_local_global(features, class_embeddings[-1])
logits = tf.ragged.map_flat_values(
core_layers.Dense(output_spec.shape[-1]), features)
valid_classes_mask = inputs.get('valid_classes_mask')
if valid_classes_mask is not None:
row_lengths = utils.diff(logits.row_splits)
valid_classes_mask = b.as_batched_model_input(valid_classes_mask)
valid_classes_mask = repeat(valid_classes_mask, row_lengths, axis=0)
def flat_fn(flat_logits):
neg_inf = tf.keras.layers.Lambda(_neg_inf_like)(flat_logits)
return utils.lambda_call(tf.where, valid_classes_mask, flat_logits,
neg_inf)
logits = tf.ragged.map_flat_values(flat_fn, logits)
return logits
def testRepeat(self, data, repeats, expected, axis=None):
result = ragged_util.repeat(data, repeats, axis)
with self.test_session():
self.assertEqual(result.eval().tolist(), expected)
