17 lines
537 B
Python
17 lines
537 B
Python
import torch
|
|
|
|
def make_coord(shape, ranges=None, flatten=True, device="cpu"):
|
|
# Make coordinates at grid centers.
|
|
coord_seqs = []
|
|
for i, n in enumerate(shape):
|
|
if ranges is None:
|
|
v0, v1 = -1, 1
|
|
else:
|
|
v0, v1 = ranges[i]
|
|
r = (v1 - v0) / (2 * n)
|
|
seq = v0 + r + (2 * r) * torch.arange(n, device=device).float()
|
|
coord_seqs.append(seq)
|
|
ret = torch.stack(torch.meshgrid(*coord_seqs), dim=-1)
|
|
if flatten:
|
|
ret = ret.view(-1, ret.shape[-1])
|
|
return ret |