gather_nd implementation with batch support.
(params, indices, batch_dims, name=None)
| 4301 | |
| 4302 | |
| 4303 | def batch_gather_nd(params, indices, batch_dims, name=None): |
| 4304 | """gather_nd implementation with batch support.""" |
| 4305 | with ops.name_scope(name, "BatchGatherND", [params, indices]): |
| 4306 | indices = ops.convert_to_tensor(indices, name="indices") |
| 4307 | params = ops.convert_to_tensor(params, name="params") |
| 4308 | |
| 4309 | if not isinstance(batch_dims, int): |
| 4310 | raise TypeError("batch_dims must be an int; got %r" % (batch_dims,)) |
| 4311 | if batch_dims < 0: |
| 4312 | raise ValueError("tf.gather_nd does not allow negative batch_dims.") |
| 4313 | params_ndims = params.shape.ndims |
| 4314 | indices_ndims = indices.shape.ndims |
| 4315 | if indices_ndims is not None and batch_dims >= indices_ndims: |
| 4316 | raise ValueError("batch_dims = %d must be less than rank(indices) = %d" % |
| 4317 | (batch_dims, indices_ndims)) |
| 4318 | if params_ndims is not None and batch_dims >= params_ndims: |
| 4319 | raise ValueError("batch_dims = %d must be less than rank(params) = %d" % |
| 4320 | (batch_dims, params_ndims)) |
| 4321 | |
| 4322 | expand = batch_dims == 0 |
| 4323 | if expand: |
| 4324 | # Normally gather_nd will be called when batch_dims == 0. |
| 4325 | # But if this function is called with batch_dims = 0, e.g. for testing |
| 4326 | # purposes, this adds a dummy batch dimension to make batch_dims = 1. |
| 4327 | params = expand_dims(params, axis=0) |
| 4328 | indices = expand_dims(indices, axis=0) |
| 4329 | batch_dims = 1 |
| 4330 | |
| 4331 | params_shape = shape(params) |
| 4332 | indices_shape = shape(indices) |
| 4333 | batch_shape = params_shape[:batch_dims] |
| 4334 | batch_size = gen_math_ops.prod(batch_shape, [0]) |
| 4335 | index_internal_ndims = rank(indices) - batch_dims - 1 |
| 4336 | indices_internal_shape = indices_shape[batch_dims:-1] |
| 4337 | |
| 4338 | # Assuming a 'params' with shape [b1, ..., bM, g1, ..., gN] and an 'indices' |
| 4339 | # with shape [b1, ..., bM, i1, ..., iK, C], where C <= N, we need to modify |
| 4340 | # 'indices' s.t. it has shape [i1, ..., iK, D], where D <= M + N and slices |
| 4341 | # to the entire 'params' tensor. |
| 4342 | # Assuming we have a batch of shape [B1, B2], we use meshgrid to create a |
| 4343 | # grid of size B1 x B2. |
| 4344 | batch_dim_list = unstack(batch_shape, axis=0) |
| 4345 | dim_ranges = [ |
| 4346 | gen_math_ops.cast(gen_math_ops._range(0, x, 1), indices.dtype) |
| 4347 | for x in batch_dim_list |
| 4348 | ] |
| 4349 | mesh_list = meshgrid(*dim_ranges, indexing="ij") if dim_ranges else [] |
| 4350 | # Then we flatten and stack the tensors to form a (B1.B2) by 2 matrix. |
| 4351 | flat_list = [reshape(x, shape=(-1,)) for x in mesh_list] |
| 4352 | index_grid = transpose(stack(flat_list, axis=0)) |
| 4353 | # We need to concatenate these batch coordinates with the internal indices. |
| 4354 | # concat -> index_grid [B1.B2, 2] with indices [i1, ..., iK, C] |
| 4355 | # So we reshape them both to [(B1.B2), i1, ..., iK, *] |
| 4356 | index_grid_shape = shape(index_grid) |
| 4357 | index_grid = reshape( |
| 4358 | index_grid, |
| 4359 | concat([ |
| 4360 | index_grid_shape[:1], |
no test coverage detected