def loop_fn(i):
inputs_i = array_ops.gather(sorted_inputs, i)
return [
array_ops.searchsorted(
inputs_i, values, out_type=dtypes.int32,
side="left"), # creates LowerBound op.
array_ops.searchsorted(
inputs_i, values, out_type=dtypes.int64, side="right")
] # creates UpperBound op.
def _PruneSparseTensor(unpruned, pruned_pattern):
"""Helper function to prune COO sparse tensor.
Given two sparse tensors 'unpruned' and 'pruned_pattern', generates another
sparse tensor with indices and values fron 'unpruned' only if its indices also
occur in pruned_pattern.
Args:
unpruned: COO matrix with unpruned indices
pruned_pattern: COO matrix with pruned pattern.
TODO(tabakg): This is far from optimal. Consider a C++ implementation.
Returns:
Indices, values, and dense_shape of the pruned matrix.
"""
pruned_indices = sparse_ops.sparse_reshape(pruned_pattern,
shape=(-1, )).indices[..., 0]
unpruned_indices = sparse_ops.sparse_reshape(unpruned,
shape=(-1, )).indices[..., 0]
best_match = array_ops.searchsorted(unpruned_indices, pruned_indices)
keep_indices = array_ops.gather(
best_match,
array_ops.where(
math_ops.equal(array_ops.gather(unpruned_indices, best_match),
pruned_indices)))
return (array_ops.gather_nd(unpruned.indices, keep_indices),
array_ops.gather_nd(unpruned.values,
keep_indices), pruned_pattern.dense_shape)
def _test2DExample(self, dtype, side, sorted_sequence, values, correct_ans):
with self.session() as session:
with self.test_scope():
tf_ans = array_ops.searchsorted(sorted_sequence, values, side=side)
tf_out = session.run(tf_ans)
self.assertAllEqual(correct_ans, tf_out)
def test1D(self):
# Test against NumPy implementation (which is 1D only).
np.random.seed(1)
for side in ['left', 'right']:
for dtype in [np.float32, np.int32]:
values = np.random.uniform(
low=-1000, high=1000, size=(10,)).astype(dtype)
unsorted = np.random.uniform(
low=-1000, high=1000, size=(20,)).astype(dtype)
sorted_sequence = np.sort(unsorted)
np_ans = np.searchsorted(sorted_sequence, values, side=side)
with self.session() as session:
with self.test_scope():
tf_ans = array_ops.searchsorted(sorted_sequence, values, side=side)
tf_out = session.run(tf_ans)
self.assertAllEqual(np_ans, tf_out)
def loop_fn(i):
inputs_i = array_ops.gather(sorted_inputs, i)
return [array_ops.searchsorted(inputs_i, values, out_type=dtypes.int32,
side="left"), # creates LowerBound op.
array_ops.searchsorted(inputs_i, values, out_type=dtypes.int64,
side="right")] # creates UpperBound op.
