def generate_shifted_anchors(
    im_h,
    im_w,
    feat_stride,
    scales,
    ratios,
    batch_size,
    dtype,
):
    """
    Enumerate all shifted anchors
    """
    height, width = np.ceil(im_h / feat_stride), np.ceil(im_w / feat_stride)

    shift_x = np.arange(0, width) * feat_stride
    shift_y = np.arange(0, height) * feat_stride
    shift_x, shift_y = np.meshgrid(shift_x, shift_y)
    shifts = np.vstack(
        (shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel())
    ).transpose()

    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    # reshape to (K*A, 4) shifted anchors

    anchors = generate_anchors(scales=scales, ratios=ratios)
    num_anchors = anchors.shape[0]
    num_shifts = shifts.shape[0]
    anchors = anchors.reshape((1, num_anchors, 4)) + shifts.reshape(
        (1, num_shifts, 4)
    ).transpose((1, 0, 2))
    anchors = anchors.reshape((1, num_shifts * num_anchors, 4))
    exp_anchors = np.repeat(anchors, repeats=batch_size, axis=0).reshape(-1, 4)
    return exp_anchors.astype(dtype)
Enumerate all shifted anchors
python/aitemplate/frontend/nn/proposal.py