utils/layout.py

from copy import deepcopy
from typing import Dict, Any


def layout_to_in_out_slice(layout, in_len, out_len=None):
    t_axis = layout.find("T")
    num_axes = len(layout)
    in_slice = [slice(None, None), ] * num_axes
    out_slice = deepcopy(in_slice)
    in_slice[t_axis] = slice(None, in_len)
    if out_len is None:
        out_slice[t_axis] = slice(in_len, None)
    else:
        out_slice[t_axis] = slice(in_len, in_len + out_len)
    return in_slice, out_slice

def step_layout_to_in_out_slice(
    layout, 
    in_len, in_step:int=1,
    out_len=None, out_step:int=1,
    in_out_diff:int=1
):
    t_axis = layout.find("T")
    num_axes = len(layout)
    in_slice = [slice(None, None), ] * num_axes
    out_slice = deepcopy(in_slice)
    
    
    in_slice[t_axis] = slice(None, in_len*in_step, in_step)
    out_start = in_len * in_step + in_out_diff - in_step
    if out_len is None:
        out_slice[t_axis] = slice(out_start, None, out_step)
    else:
        out_slice[t_axis] = slice(out_start, out_start + out_len*out_step,out_step)
    return in_slice, out_slice


def parse_layout_shape(layout: str) -> Dict[str, Any]:
    r"""

    Parameters
    ----------
    layout: str
            e.g., "NTHWC", "NHWC".

    Returns
    -------
    ret:    Dict
    """
    batch_axis = layout.find("N")
    t_axis = layout.find("T")
    h_axis = layout.find("H")
    w_axis = layout.find("W")
    c_axis = layout.find("C")
    return {
        "batch_axis": batch_axis,
        "t_axis": t_axis,
        "h_axis": h_axis,
        "w_axis": w_axis,
        "c_axis": c_axis,
    }