(inp: HLOTensor, size: Sequence[int])
| 47 | |
| 48 | |
| 49 | def resize_nearest(inp: HLOTensor, size: Sequence[int]): |
| 50 | N, C, IH, IW = inp.shape |
| 51 | OH, OW = size |
| 52 | assert IH != 0 and IW != 0 and OH != 0 and OW != 0, f"{inp.shape}, {size}" |
| 53 | |
| 54 | ihidx = resize_nearest_helper(IH, OH) |
| 55 | iwidx = resize_nearest_helper(IW, OW) |
| 56 | ihidx = ihidx.reshape((OH, 1)).broadcast_to((OH, OW)) |
| 57 | iwidx = iwidx.reshape((1, OW)).broadcast_to((OH, OW)) |
| 58 | iidx = ihidx * IW + iwidx |
| 59 | iidx = iidx.broadcast_to((N, C, OH, OW)).reshape((N, C, OH * OW)) |
| 60 | |
| 61 | dim0 = ( |
| 62 | iota(dtype=np.int32, shape=(N,), dimension=-1) |
| 63 | .reshape((N, 1, 1)) |
| 64 | .broadcast_to(iidx.shape) |
| 65 | ) |
| 66 | dim1 = ( |
| 67 | iota(dtype=np.int32, shape=(C,), dimension=-1) |
| 68 | .reshape((1, C, 1)) |
| 69 | .broadcast_to(iidx.shape) |
| 70 | ) |
| 71 | iidx = stack([dim0, dim1, iidx], axis=-1) |
| 72 | |
| 73 | inp = inp.reshape((N, C, -1)) |
| 74 | out = xla_gather( |
| 75 | inp, |
| 76 | iidx, |
| 77 | slice_sizes=(1, 1, 1), |
| 78 | offset_dims=tuple(), |
| 79 | collapsed_slice_dims=(0, 1, 2), |
| 80 | start_index_map=(0, 1, 2), |
| 81 | ) |
| 82 | return out.reshape((N, C, OH, OW)) |
| 83 | |
| 84 | |
| 85 | def resize_linear(inp: HLOTensor, size: Sequence[int]): |
no test coverage detected