def mesh_grid(B, H, W):
# mesh grid
x_base = F.arange(0, W)
x_base = F.tile(x_base, (B, H, 1))
y_base = F.arange(0, H) # BHW
y_base = F.tile(y_base, (B, W, 1)).transpose(0, 2, 1)
base_grid = F.stack([x_base, y_base], 1) # B2HW
return base_grid
def mesh_grid_mge(B, H, W):
# mesh grid
x_base = F.arange(0, W)
x_base = F.tile(x_base, (B, H, 1))
y_base = F.arange(0, H) # BHW
y_base = F.tile(y_base, (B, W, 1)).transpose(0, 2, 1)
ones = F.ones_like(x_base)
base_grid = F.stack([x_base, y_base, ones], 1) # B3HW
return base_grid
def tile_func(inp):
return F.tile(inp=inp, reps=reps)
[(64, 512, 16, 16)],
True,
1000,
),
(
"topk",
lambda x: MF.topk(x, 10),
lambda x: torch.topk(x, 10),
[(100, 100)],
[(1000, 1000)],
True,
1000,
),
(
"tile",
lambda x: MF.tile(x, (2, ) * len(x.shape)),
lambda x: torch.tile(x, (2, ) * len(x.shape)),
[(100, 100)],
[(64, 512, 16, 16)],
True,
1000,
),
(
"transpose",
lambda x: MF.transpose(x,
list(range(len(x.shape)))[::-1]),
lambda x: torch.permute(x,
list(range(len(x.shape)))[::-1]),
[(100, 100)],
[(64, 512, 16, 16)],
True,