pytorch-image-models

1
from typing import Tuple
2

3
import torch
4

5

6
def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]:
7
    """generate N-D grid in dimension order.
8

9
    The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.
10

11
    That is, the statement
12
    [X1,X2,X3] = ndgrid(x1,x2,x3)
13

14
    produces the same result as
15

16
    [X2,X1,X3] = meshgrid(x2,x1,x3)
17

18
    This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make
19
    torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy').
20

21
    """
22
    try:
23
        return torch.meshgrid(*tensors, indexing='ij')
24
    except TypeError:
25
        # old PyTorch < 1.10 will follow this path as it does not have indexing arg,
26
        # the old behaviour of meshgrid was 'ij'
27
        return torch.meshgrid(*tensors)
28

29

30
def meshgrid(*tensors) -> Tuple[torch.Tensor, ...]:
31
    """generate N-D grid in spatial dim order.
32

33
    The meshgrid function is similar to ndgrid except that the order of the
34
    first two input and output arguments is switched.
35

36
    That is, the statement
37

38
    [X,Y,Z] = meshgrid(x,y,z)
39
    produces the same result as
40

41
    [Y,X,Z] = ndgrid(y,x,z)
42
    Because of this, meshgrid is better suited to problems in two- or three-dimensional Cartesian space,
43
    while ndgrid is better suited to multidimensional problems that aren't spatially based.
44
    """
45

46
    # NOTE: this will throw in PyTorch < 1.10 as meshgrid did not support indexing arg or have
47
    # capability of generating grid in xy order before then.
48
    return torch.meshgrid(*tensors, indexing='xy')
49

50