r"""Gather slices from params according to indices with leading batch dims. This operation assumes that the leading `batch_dims` dimensions of `indices` and `params` are batch dimensions; and performs a `tf.gather` operation within each batch. (If `batch_dims` is not specified, then it defaul
(params, indices, batch_dims, axis=None)
| 4003 | |
| 4004 | |
| 4005 | def _batch_gather(params, indices, batch_dims, axis=None): |
| 4006 | r"""Gather slices from params according to indices with leading batch dims. |
| 4007 | |
| 4008 | This operation assumes that the leading `batch_dims` dimensions of `indices` |
| 4009 | and `params` are batch dimensions; and performs a `tf.gather` operation within |
| 4010 | each batch. (If `batch_dims` is not specified, then it defaults to |
| 4011 | `rank(indices)-1`.) In the case in which `batch_dims==0`, this operation |
| 4012 | is equivalent to `tf.gather`. |
| 4013 | |
| 4014 | Args: |
| 4015 | params: A Tensor. The tensor from which to gather values. |
| 4016 | indices: A Tensor. Must be one of the following types: int32, int64. Index |
| 4017 | tensor. Must be in range `[0, params.shape[batch_dims]]`. |
| 4018 | batch_dims: An integer or none. The number of batch dimensions. Must be |
| 4019 | less than `rank(indices)`. Defaults to `rank(indices) - 1` if None. |
| 4020 | axis: A `Tensor`. Must be one of the following types: `int32`, `int64`. The |
| 4021 | `axis` in `params` to gather `indices` from. Must be greater than or equal |
| 4022 | to `batch_dims`. Defaults to the first non-batch dimension. Supports |
| 4023 | negative indexes. |
| 4024 | |
| 4025 | Returns: |
| 4026 | A Tensor. Has the same type as `params`. |
| 4027 | |
| 4028 | Raises: |
| 4029 | ValueError: if `indices` has an unknown shape. |
| 4030 | """ |
| 4031 | if batch_dims is not None and not isinstance(batch_dims, int): |
| 4032 | raise TypeError("batch_dims must be an int; got %r" % (batch_dims,)) |
| 4033 | indices = ops.convert_to_tensor(indices, name="indices") |
| 4034 | params = ops.convert_to_tensor(params, name="params") |
| 4035 | |
| 4036 | indices_ndims = indices.shape.ndims |
| 4037 | if indices_ndims is None: |
| 4038 | raise ValueError("tf.gather does not allow indices with unknown " |
| 4039 | "rank when batch_dims is specified.") |
| 4040 | if batch_dims is None: |
| 4041 | batch_dims = indices_ndims - 1 |
| 4042 | if batch_dims < 0: |
| 4043 | batch_dims += indices_ndims |
| 4044 | if batch_dims < 0 or batch_dims >= indices_ndims: |
| 4045 | raise ValueError("batch_dims = %d must be less than rank(indices) = %d" % |
| 4046 | (batch_dims, indices_ndims)) |
| 4047 | if params.shape.ndims is not None and batch_dims >= params.shape.ndims: |
| 4048 | raise ValueError("batch_dims = %d must be less than rank(params) = %d" % |
| 4049 | (batch_dims, params.shape.ndims)) |
| 4050 | |
| 4051 | # Handle axis by transposing the axis dimension to be the first non-batch |
| 4052 | # dimension, recursively calling batch_gather with axis=0, and then |
| 4053 | # transposing the result to put the pre-axis dimensions before the indices |
| 4054 | # dimensions. |
| 4055 | if axis is not None and axis != batch_dims: |
| 4056 | # Adjust axis to be positive. |
| 4057 | if not isinstance(axis, int): |
| 4058 | axis = tf.where(axis < 0, axis + array_ops.rank(params), axis) |
| 4059 | elif axis < 0 and params.shape.ndims is None: |
| 4060 | axis = axis + array_ops.rank(params) |
| 4061 | else: |
| 4062 | if (axis < -params.shape.ndims) or (axis >= params.shape.ndims): |
no test coverage detected