Helper function for unsorted segment ops. Gathers params for positive segment ids and gathers 0 for inputs with negative segment id. Also returns the clipped indices and a boolean mask with the same shape as ids where a positive id is masked as true. With this, the latter two
(params,
ids,
zero_clipped_indices=None,
is_positive=None)
| 396 | |
| 397 | |
| 398 | def _GatherDropNegatives(params, |
| 399 | ids, |
| 400 | zero_clipped_indices=None, |
| 401 | is_positive=None): |
| 402 | """ Helper function for unsorted segment ops. |
| 403 | |
| 404 | Gathers params for |
| 405 | positive segment ids and gathers 0 for inputs with negative segment id. |
| 406 | Also returns the clipped indices and a boolean mask with the same shape |
| 407 | as ids where a positive id is masked as true. With this, the latter two |
| 408 | can be passed as arguments to this function to reuse them. |
| 409 | """ |
| 410 | if zero_clipped_indices is None: |
| 411 | zero_clipped_indices = math_ops.maximum(ids, array_ops.zeros_like(ids)) |
| 412 | gathered = array_ops.gather(params, zero_clipped_indices) |
| 413 | if is_positive is None: |
| 414 | is_positive = math_ops.greater_equal(ids, 0) |
| 415 | # tf.where(condition, x, y) requires condition to have the same shape as x |
| 416 | # and y. |
| 417 | # todo(philjd): remove this if tf.where supports broadcasting (#9284) |
| 418 | for _ in range(gathered.shape.ndims - is_positive.shape.ndims): |
| 419 | is_positive = array_ops.expand_dims(is_positive, -1) |
| 420 | is_positive = ( |
| 421 | is_positive & array_ops.ones_like(gathered, dtype=dtypes.bool)) |
| 422 | # replace gathered params of negative indices with 0 |
| 423 | zero_slice = array_ops.zeros_like(gathered) |
| 424 | return (array_ops.where(is_positive, gathered, zero_slice), |
| 425 | zero_clipped_indices, is_positive) |
| 426 | |
| 427 | |
| 428 | def _UnsortedSegmentMinOrMaxGrad(op, grad): |
no test coverage detected