-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
search_sorted.ts
116 lines (109 loc) · 4.67 KB
/
search_sorted.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
/**
* @license
* Copyright 2022 Google LLC. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
import {ENGINE} from '../engine';
import {SearchSorted, SearchSortedAttrs, SearchSortedInputs} from '../kernel_names';
import {Tensor} from '../tensor';
import {convertToTensor} from '../tensor_util_env';
import {TensorLike} from '../types';
import {sizeFromShape} from '../util_base';
import {op} from './operation';
import {reshape} from './reshape';
const INT32_MAX = 2147483648;
/**
* Searches for where a value would go in a sorted sequence.
*
* This is not a method for checking containment (like javascript in).
*
* The typical use case for this operation is "binning", "bucketing", or
* "discretizing". The values are assigned to bucket-indices based on the edges
* listed in 'sortedSequence'. This operation returns the bucket-index for each
* value.
*
* The side argument controls which index is returned if a value lands exactly
* on an edge.
*
* The axis is not settable for this operation. It always operates on the
* innermost dimension (axis=-1). The operation will accept any number of outer
* dimensions.
*
* Note: This operation assumes that 'sortedSequence' is sorted along the
* innermost axis, maybe using 'sort(..., axis=-1)'. If the sequence is not
* sorted no error is raised and the content of the returned tensor is not well
* defined.
*
* ```js
* const edges = tf.tensor1d([-1, 3.3, 9.1, 10.0]);
* let values = tf.tensor1d([0.0, 4.1, 12.0]);
* const result1 = tf.searchSorted(edges, values, 'left');
* result1.print(); // [1, 2, 4]
*
* const seq = tf.tensor1d([0, 3, 9, 10, 10]);
* values = tf.tensor1d([0, 4, 10]);
* const result2 = tf.searchSorted(seq, values, 'left');
* result2.print(); // [0, 2, 3]
* const result3 = tf.searchSorted(seq, values, 'right');
* result3.print(); // [1, 2, 5]
*
* const sortedSequence = tf.tensor2d([[0., 3., 8., 9., 10.],
* [1., 2., 3., 4., 5.]]);
* values = tf.tensor2d([[9.8, 2.1, 4.3],
* [0.1, 6.6, 4.5, ]]);
* const result4 = tf.searchSorted(sortedSequence, values, 'left');
* result4.print(); // [[4, 1, 2], [0, 5, 4]]
* ```
* @param sortedSequence: N-D. Sorted sequence.
* @param values: N-D. Search values.
* @param side: 'left'|'right'. Defaults to 'left'. 'left' corresponds to lower
* bound and 'right' to upper bound.
* @return An N-D int32 tensor the size of values containing the result of
* applying either lower bound or upper bound (depending on side) to each
* value. The result is not a global index to the entire Tensor, but the
* index in the last dimension.
* @doc {heading: 'Operations', subheading: 'Evaluation'}
*/
function searchSorted_(
sortedSequence: Tensor|TensorLike, values: Tensor|TensorLike,
side: 'left'|'right' = 'left'): Tensor {
const $sortedSequence =
convertToTensor(sortedSequence, 'sortedSequence', 'searchSorted');
const $values = convertToTensor(values, 'values', 'searchSorted');
const sequenceSize = $sortedSequence.shape[$sortedSequence.shape.length - 1];
const valuesSize = $values.shape[$values.shape.length - 1];
const $sortedSequence2D = reshape($sortedSequence, [-1, sequenceSize]);
const $values2D = reshape($values, [-1, valuesSize]);
if ($sortedSequence2D.rank < 2) {
throw new Error(`Sorted input argument must be at least 2-dimensional`);
}
if ($sortedSequence2D.shape[0] !== $values2D.shape[0]) {
throw new Error(
`Leading dimension of 'sortedSequence' and 'values' must match.`);
}
if (sizeFromShape($values2D.shape) >= INT32_MAX) {
throw new Error(`values tensor size must less than ${INT32_MAX}`);
}
if ($sortedSequence2D.shape[1] >= INT32_MAX) {
throw new Error(`trailing dim_size must less than ${
INT32_MAX} for int32 output type, was ${$sortedSequence2D.shape[1]}`);
}
const inputs: SearchSortedInputs = {
sortedSequence: $sortedSequence2D,
values: $values2D,
};
const attrs: SearchSortedAttrs = {side};
return ENGINE.runKernel(SearchSorted, inputs as {}, attrs as {});
}
export const searchSorted = /* @__PURE__ */ op({searchSorted_});