def get_num_rightmost_static_elements(shape, num_rightmost_dims: int = None) -> int:
    """
    Returns number of elements in rightmost max contiguous static dimensions.
    If the rightmost dim is dynamic, returns 1.

    If num_rightmost_dims is specified, only look into num_rightmost_dims.
    Otherwise, look into all dims.
    If num_rightmost_dims == 0, returns 1.

    This is useful when calculating alignment.

    e.g.
    shape = [IntImm(2), IntImm(4)] returns 2 * 4 = 8.
    shape = [IntImm(2), IntVar(4, 8)] returns 1.
    shape = [IntImm(9), IntVar(4, 8), IntImm(3)] returns 3.

    shape = [IntImm(2), IntImm(4), IntImm(3)], num_rightmost_dims = None returns 24.
    shape = [IntImm(2), IntImm(4), IntImm(3)], num_rightmost_dims = 1 returns 3.
    shape = [IntImm(2), IntImm(4), IntImm(3)], num_rightmost_dims = 0 returns 1.
    """
    from aitemplate.compiler.base import IntImm

    res = 1

    for idx, dim in enumerate(reversed(shape)):
        if num_rightmost_dims is not None and idx >= num_rightmost_dims:
            break
        if not isinstance(dim, IntImm):
            break
        res *= dim._attrs["values"][0]
    return res

    Returns number of elements in rightmost max contiguous static dimensions.
    If the rightmost dim is dynamic, returns 1.

    If num_rightmost_dims is specified, only look into num_rightmost_dims.
    Otherwise, look into all dims.
    If num_rightmost_dims == 0, returns 1.

    This is useful when calculating alignment.

    e.g.
    shape = [IntImm(2), IntImm(4)] returns 2 * 4 = 8.
    shape = [IntImm(2), IntVar(4, 8)] returns 1.
    shape = [IntImm(9), IntVar(4, 8), IntImm(3)] returns 3.

    shape = [IntImm(2), IntImm(4), IntImm(3)], num_rightmost_dims = None returns 24.
    shape = [IntImm(2), IntImm(4), IntImm(3)], num_rightmost_dims = 1 returns 3.
    shape = [IntImm(2), IntImm(4), IntImm(3)], num_rightmost_dims = 0 returns 1.
    
python/aitemplate/utils/shape_utils.py