Filter all but `selected_id` out of `ids`. Args: ids: `int64` `Tensor` or `SparseTensor` of IDs. selected_id: Int id to select. Returns: `SparseTensor` of same dimensions as `ids`. This contains only the entries equal to `selected_id`.
(ids, selected_id)
| 2224 | |
| 2225 | |
| 2226 | def _select_class_id(ids, selected_id): |
| 2227 | """Filter all but `selected_id` out of `ids`. |
| 2228 | |
| 2229 | Args: |
| 2230 | ids: `int64` `Tensor` or `SparseTensor` of IDs. |
| 2231 | selected_id: Int id to select. |
| 2232 | |
| 2233 | Returns: |
| 2234 | `SparseTensor` of same dimensions as `ids`. This contains only the entries |
| 2235 | equal to `selected_id`. |
| 2236 | """ |
| 2237 | ids = sparse_tensor.convert_to_tensor_or_sparse_tensor(ids) |
| 2238 | if isinstance(ids, sparse_tensor.SparseTensor): |
| 2239 | return sparse_ops.sparse_retain(ids, math_ops.equal(ids.values, |
| 2240 | selected_id)) |
| 2241 | |
| 2242 | # TODO(ptucker): Make this more efficient, maybe add a sparse version of |
| 2243 | # tf.equal and tf.reduce_any? |
| 2244 | |
| 2245 | # Shape of filled IDs is the same as `ids` with the last dim collapsed to 1. |
| 2246 | ids_shape = array_ops.shape(ids, out_type=dtypes.int64) |
| 2247 | ids_last_dim = array_ops.size(ids_shape) - 1 |
| 2248 | filled_selected_id_shape = math_ops.reduced_shape(ids_shape, |
| 2249 | array_ops.reshape( |
| 2250 | ids_last_dim, [1])) |
| 2251 | |
| 2252 | # Intersect `ids` with the selected ID. |
| 2253 | filled_selected_id = array_ops.fill(filled_selected_id_shape, |
| 2254 | math_ops.cast(selected_id, dtypes.int64)) |
| 2255 | result = sets.set_intersection(filled_selected_id, ids) |
| 2256 | return sparse_tensor.SparseTensor( |
| 2257 | indices=result.indices, values=result.values, dense_shape=ids_shape) |
| 2258 | |
| 2259 | |
| 2260 | def _maybe_select_class_id(labels, predictions_idx, selected_id=None): |