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.