r"""Repeat elements of an array. Args: inp: input tensor. repeats: the number of repetitions for each element. axis: the axis along which to repeat values. By default, use the flattened input array, and return a flat output array. Returns: output
(inp: Tensor, repeats: int, axis: Optional[int] = None)
| 1088 | |
| 1089 | |
| 1090 | def repeat(inp: Tensor, repeats: int, axis: Optional[int] = None): |
| 1091 | r"""Repeat elements of an array. |
| 1092 | |
| 1093 | Args: |
| 1094 | inp: input tensor. |
| 1095 | repeats: the number of repetitions for each element. |
| 1096 | axis: the axis along which to repeat values. By default, use the |
| 1097 | flattened input array, and return a flat output array. |
| 1098 | |
| 1099 | Returns: |
| 1100 | output tensor. |
| 1101 | |
| 1102 | Examples: |
| 1103 | >>> import numpy as np |
| 1104 | >>> x = Tensor([[1, 2], [3, 4]], np.int32) |
| 1105 | >>> F.repeat(x, 2, axis=0) |
| 1106 | Tensor([[1 2] |
| 1107 | [1 2] |
| 1108 | [3 4] |
| 1109 | [3 4]], dtype=int32, device=xpux:0) |
| 1110 | """ |
| 1111 | if axis is None: |
| 1112 | inp = inp.reshape(-1) # flatten |
| 1113 | axis = 0 |
| 1114 | shape = astensor1d(inp.shape, inp, dtype="int32", device=inp.device) |
| 1115 | # assume inp.ndim is not changed during trace |
| 1116 | max_axis = len(shape) - 1 |
| 1117 | assert axis >= 0 and axis <= max_axis |
| 1118 | assert repeats >= 1 |
| 1119 | |
| 1120 | base_shape, bcast_shape, target_shape = [], [], [] |
| 1121 | if axis != 0: |
| 1122 | target_shape.append(shape[:axis]) |
| 1123 | base_shape.extend([shape[: axis + 1], [1,]]) |
| 1124 | bcast_shape.extend([shape[: axis + 1], [repeats,]]) |
| 1125 | target_shape.extend( |
| 1126 | [shape[axis] * repeats,] |
| 1127 | ) |
| 1128 | if axis + 1 <= max_axis: |
| 1129 | base_shape.append(shape[axis + 1 :]) |
| 1130 | bcast_shape.append(shape[axis + 1 :]) |
| 1131 | target_shape.append(shape[axis + 1 :]) |
| 1132 | |
| 1133 | base_shape = astensor1d(base_shape) |
| 1134 | bcast_shape = astensor1d(bcast_shape) |
| 1135 | target_shape = astensor1d(target_shape) |
| 1136 | out = broadcast_to(inp.reshape(base_shape), bcast_shape).reshape(target_shape) |
| 1137 | return out |
| 1138 | |
| 1139 | |
| 1140 | def _tile_one_dim(inp, rep, axis): |
nothing calls this directly
no test coverage detected