Slice out a subset of the tensor. Args: labeled_tensor: The input tensor. selection: A dictionary mapping an axis name to a scalar, slice or list of values to select. Currently supports two types of selections: (a) Any number of scalar and/or slice selections. (b) Ex
(labeled_tensor, selection, name=None)
| 55 | tc.Union(slice, collections_abc.Hashable, list)), |
| 56 | tc.Optional(string_types)) |
| 57 | def select(labeled_tensor, selection, name=None): |
| 58 | """Slice out a subset of the tensor. |
| 59 | |
| 60 | Args: |
| 61 | labeled_tensor: The input tensor. |
| 62 | selection: A dictionary mapping an axis name to a scalar, slice or list of |
| 63 | values to select. Currently supports two types of selections: |
| 64 | (a) Any number of scalar and/or slice selections. |
| 65 | (b) Exactly one list selection, without any scalars or slices. |
| 66 | name: Optional op name. |
| 67 | |
| 68 | Returns: |
| 69 | The selection as a `LabeledTensor`. |
| 70 | |
| 71 | Raises: |
| 72 | ValueError: If the tensor doesn't have an axis in the selection or if |
| 73 | that axis lacks labels. |
| 74 | KeyError: If any labels in a selection are not found in the original axis. |
| 75 | NotImplementedError: If you attempt to combine a list selection with |
| 76 | scalar selection or another list selection. |
| 77 | """ |
| 78 | with ops.name_scope(name, 'lt_select', [labeled_tensor]) as scope: |
| 79 | labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor) |
| 80 | |
| 81 | slices = {} |
| 82 | indexers = {} |
| 83 | for axis_name, value in selection.items(): |
| 84 | if axis_name not in labeled_tensor.axes: |
| 85 | raise ValueError( |
| 86 | 'The tensor does not have an axis named %s. Its axes are: %r' % |
| 87 | (axis_name, labeled_tensor.axes.keys())) |
| 88 | axis = labeled_tensor.axes[axis_name] |
| 89 | if axis.labels is None: |
| 90 | raise ValueError( |
| 91 | 'The axis named %s does not have labels. The axis is: %r' % |
| 92 | (axis_name, axis)) |
| 93 | |
| 94 | if isinstance(value, slice): |
| 95 | # TODO(shoyer): consider deprecating using slices in favor of lists |
| 96 | if value.start is None: |
| 97 | start = None |
| 98 | else: |
| 99 | start = axis.index(value.start) |
| 100 | |
| 101 | if value.stop is None: |
| 102 | stop = None |
| 103 | else: |
| 104 | # For now, follow the pandas convention of making labeled slices |
| 105 | # inclusive of both bounds. |
| 106 | stop = axis.index(value.stop) + 1 |
| 107 | |
| 108 | if value.step is not None: |
| 109 | raise NotImplementedError('slicing with a step is not yet supported') |
| 110 | |
| 111 | slices[axis_name] = slice(start, stop) |
| 112 | |
| 113 | # Needs to be after checking for slices, since slice objects claim to be |
| 114 | # instances of collections_abc.Hashable but hash() on them fails. |
no test coverage detected