(inp: HLOTensor, size: Sequence[int])
| 83 | |
| 84 | |
| 85 | def resize_linear(inp: HLOTensor, size: Sequence[int]): |
| 86 | N, C, IH, IW = inp.shape |
| 87 | OH, OW = size |
| 88 | assert IH != 0 and IW != 0 and OH != 0 and OW != 0, f"{inp.shape}, {size}" |
| 89 | |
| 90 | ah0s, ih0s, ah1s, ih1s = resize_linear_helper(IH, OH) |
| 91 | aw0s, iw0s, aw1s, iw1s = resize_linear_helper(IW, OW) |
| 92 | |
| 93 | ih0s = ih0s.reshape((OH, 1)).broadcast_to((OH, OW)) |
| 94 | ih1s = ih1s.reshape((OH, 1)).broadcast_to((OH, OW)) |
| 95 | iw0s = iw0s.reshape((1, OW)).broadcast_to((OH, OW)) |
| 96 | iw1s = iw1s.reshape((1, OW)).broadcast_to((OH, OW)) |
| 97 | iidx00 = ih0s * IW + iw0s |
| 98 | iidx01 = ih0s * IW + iw1s |
| 99 | iidx10 = ih1s * IW + iw0s |
| 100 | iidx11 = ih1s * IW + iw1s |
| 101 | iidx00 = iidx00.broadcast_to((N, C, OH, OW)).reshape((N, C, -1)) |
| 102 | iidx01 = iidx01.broadcast_to((N, C, OH, OW)).reshape((N, C, -1)) |
| 103 | iidx10 = iidx10.broadcast_to((N, C, OH, OW)).reshape((N, C, -1)) |
| 104 | iidx11 = iidx11.broadcast_to((N, C, OH, OW)).reshape((N, C, -1)) |
| 105 | |
| 106 | dim0 = ( |
| 107 | iota(dtype=np.int32, shape=(N,), dimension=-1) |
| 108 | .reshape((N, 1, 1)) |
| 109 | .broadcast_to(iidx00.shape) |
| 110 | ) |
| 111 | dim1 = ( |
| 112 | iota(dtype=np.int32, shape=(C,), dimension=-1) |
| 113 | .reshape((1, C, 1)) |
| 114 | .broadcast_to(iidx00.shape) |
| 115 | ) |
| 116 | iidx00 = stack([dim0, dim1, iidx00], axis=-1) |
| 117 | iidx01 = stack([dim0, dim1, iidx01], axis=-1) |
| 118 | iidx10 = stack([dim0, dim1, iidx10], axis=-1) |
| 119 | iidx11 = stack([dim0, dim1, iidx11], axis=-1) |
| 120 | |
| 121 | inp = inp.reshape((N, C, -1)) |
| 122 | out00 = xla_gather( |
| 123 | inp, |
| 124 | iidx00, |
| 125 | slice_sizes=(1, 1, 1), |
| 126 | offset_dims=tuple(), |
| 127 | collapsed_slice_dims=(0, 1, 2), |
| 128 | start_index_map=(0, 1, 2), |
| 129 | ) |
| 130 | out01 = xla_gather( |
| 131 | inp, |
| 132 | iidx01, |
| 133 | slice_sizes=(1, 1, 1), |
| 134 | offset_dims=tuple(), |
| 135 | collapsed_slice_dims=(0, 1, 2), |
| 136 | start_index_map=(0, 1, 2), |
| 137 | ) |
| 138 | out10 = xla_gather( |
| 139 | inp, |
| 140 | iidx10, |
| 141 | slice_sizes=(1, 1, 1), |
| 142 | offset_dims=tuple(), |
no test coverage detected