Addds the batch offsets to the given indices and returns the results.
(params_shape, indices, batch_dims)
| 479 | |
| 480 | |
| 481 | def _GetBatchIndices(params_shape, indices, batch_dims): |
| 482 | """Addds the batch offsets to the given indices and returns the results.""" |
| 483 | batch_indices = indices |
| 484 | indices_ndims = indices.shape.ndims |
| 485 | indices_dtype = indices.dtype.base_dtype |
| 486 | casted_params_shape = math_ops.cast(params_shape, indices_dtype) |
| 487 | accum_dim_value = array_ops.ones((), dtype=indices_dtype) |
| 488 | for dim in range(batch_dims, 0, -1): |
| 489 | dim_value = casted_params_shape[dim - 1] |
| 490 | accum_dim_value *= casted_params_shape[dim] |
| 491 | start = array_ops.zeros((), dtype=indices_dtype) |
| 492 | step = array_ops.ones((), dtype=indices_dtype) |
| 493 | dim_indices = math_ops.range(start, dim_value, step) |
| 494 | dim_indices *= accum_dim_value |
| 495 | dim_shape = array_ops.stack( |
| 496 | [1] * (dim - 1) + [dim_value] + [1] * (indices_ndims - dim), axis=0) |
| 497 | batch_indices += array_ops.reshape(dim_indices, dim_shape) |
| 498 | |
| 499 | return batch_indices |
| 500 | |
| 501 | |
| 502 | def _BatchGatherGrad( |