Source code for vis4d.op.box.anchor.util

"""Anchor utils."""

from __future__ import annotations

from torch import Tensor



[docs]
def meshgrid(
    x_grid: Tensor, y_grid: Tensor, row_major: bool = True
) -> tuple[Tensor, Tensor]:
    """Generate mesh grid of x and y.

    Args:
        x_grid (Tensor): Grids of x dimension.
        y_grid (Tensor): Grids of y dimension.
        row_major (bool, optional): Whether to return y grids first.
            Defaults to True.

    Returns:
        tuple[Tensor]: The mesh grids of x and y.
    """
    # use shape instead of len to keep tracing while exporting to onnx
    xx = x_grid.repeat(y_grid.shape[0])
    yy = y_grid.view(-1, 1).repeat(1, x_grid.shape[0]).view(-1)
    if row_major:
        return xx, yy
    return yy, xx