linear_expansion

Bases: transformation

The linear data expansion function.

It performs the linear expansion of the input vector, and returns the expansion result. The class inherits from the base expansion class (i.e., the transformation class in the module directory).

...

Notes

For input vector \(\mathbf{x} \in R^m\), its linear expansion can be based on one of the following equations: $$ \begin{align} \kappa(\mathbf{x}) &= c \mathbf{x} \in {R}^D, \\ \kappa(\mathbf{x}) &= \mathbf{x} \mathbf{C}_{post} \in {R}^D,\\ \kappa(\mathbf{x}) &= \mathbf{C}_{pre} \mathbf{x} \in {R}^{D}, \end{align} $$ where \(c \in {R}\), \(\mathbf{C}_{post}, \mathbf{C}_{pre} \in {R}^{m \times m}\) denote the provided constant scalar and linear transformation matrices, respectively. Linear data expansion will not change the data vector dimensions, and the output data vector dimension \(D=m\).

By default, the input and output can also be processed with the optional pre- or post-processing functions in the linear expansion function.

Attributes:

Name	Type	Description
name	str, default = 'linear_expansion'	Name of the expansion function.

Methods:

Name	Description
__init__	It performs the initialization of the expansion function.
calculate_D	It calculates the expansion space dimension D based on the input dimension parameter m.
forward	It implements the abstract forward method declared in the base expansion class.

Source code in tinybig/expansion/basic_expansion.py

class linear_expansion(transformation):
    r"""
    The linear data expansion function.

    It performs the linear expansion of the input vector, and returns the expansion result.
    The class inherits from the base expansion class (i.e., the transformation class in the module directory).

    ...

    Notes
    ----------
    For input vector $\mathbf{x} \in R^m$, its linear expansion can be based on one of the following equations:
    $$
    \begin{align}
        \kappa(\mathbf{x}) &= c \mathbf{x} \in {R}^D, \\\\
        \kappa(\mathbf{x}) &= \mathbf{x} \mathbf{C}\_{post} \in {R}^D,\\\\
        \kappa(\mathbf{x}) &= \mathbf{C}\_{pre} \mathbf{x} \in {R}^{D},
    \end{align}
    $$
    where $c \in {R}$, $\mathbf{C}_{post}, \mathbf{C}_{pre} \in {R}^{m \times m}$ denote the provided
    constant scalar and linear transformation matrices, respectively.
    Linear data expansion will not change the data vector dimensions, and the output data vector dimension $D=m$.

    By default, the input and output can also be processed with the optional pre- or post-processing functions
    in the linear expansion function.

    Attributes
    ----------
    name: str, default = 'linear_expansion'
        Name of the expansion function.

    Methods
    ----------
    __init__
        It performs the initialization of the expansion function.

    calculate_D
        It calculates the expansion space dimension D based on the input dimension parameter m.

    forward
        It implements the abstract forward method declared in the base expansion class.

    """
    def __init__(self, name='linear_expansion', c=None, pre_C=None, post_C=None, *args, **kwargs):
        r"""
        The initialization method of the linear expansion function.

        It initializes a linear expansion object based on the input function name.
        This method will also call the initialization method of the base class as well.

        Parameters
        ----------
        name: str, default = 'linear_expansion'
            The name of the linear expansion function.
        c: float | torch.Tensor, default = None
            The scalar $c$ of the linear expansion.
        pre_C: torch.Tensor, default = None
            The $\mathbf{C}_{pre}$ matrix of the linear expansion.
        post_C: torch.Tensor, default = None
            The $\mathbf{C}_{post}$ matrix of the linear expansion.

        Returns
        ----------
        transformation
            The linear expansion function.
        """
        super().__init__(name=name, *args, **kwargs)
        self.c = c
        self.pre_C = pre_C
        self.post_C = post_C

    def calculate_D(self, m: int):
        r"""
        The expansion dimension calculation method.

        It calculates the intermediate expansion space dimension based on the input dimension parameter m.
        For the linear expansion function, the expansion space dimension equals to the input space dimension, i.e.,
        $$ D = m. $$

        Parameters
        ----------
        m: int
            The dimension of the input space.

        Returns
        -------
        int
            The dimension of the expansion space.
        """
        return m

    def forward(self, x: torch.Tensor, device='cpu', c=None, pre_C=None, post_C=None, *args, **kwargs):
        r"""
        The forward method of the data expansion function.

        It performs the linear data expansion of the input data and returns the expansion result
        according to one of the following equation:
        $$
        \begin{align}
            \kappa(\mathbf{x}) &= c \mathbf{x} \in {R}^D, \\\\
            \kappa(\mathbf{x}) &= \mathbf{x} \mathbf{C}\_{post} \in {R}^D,\\\\
            \kappa(\mathbf{x}) &= \mathbf{C}\_{pre} \mathbf{x} \in {R}^{D},
        \end{align}
        $$
        where $c \in {R}$, $\mathbf{C}_{post}, \mathbf{C}_{pre} \in {R}^{m \times m}$ denote the provided
        constant scalar and linear transformation matrices, respectively.


        Parameters
        ----------
        x: torch.Tensor
            The input data vector.
        device: str, default = 'cpu'
            The device to perform the data expansion.
        c: float | torch.Tensor, default = None
            The scalar $c$ of the linear expansion.
        pre_C: torch.Tensor, default = None
            The $\mathbf{C}_{pre}$ matrix of the linear expansion.
        post_C: torch.Tensor, default = None
            The $\mathbf{C}_{post}$ matrix of the linear expansion.

        Returns
        ----------
        torch.Tensor
            The expanded data vector of the input.
        """
        b, m = x.shape
        x = self.pre_process(x=x, device=device)

        c = c if c is not None else self.c
        pre_C = pre_C if pre_C is not None else self.pre_C
        post_C = post_C if post_C is not None else self.post_C

        if c is not None:
            expansion = c * x
        elif pre_C is not None:
            assert pre_C.size(-1) == x.size(0)
            expansion = torch.matmul(pre_C, x)
        elif post_C is not None:
            assert x.size(-1) == post_C.size(0)
            expansion = torch.matmul(x, post_C)
        else:
            expansion = x

        assert expansion.shape == (b, self.calculate_D(m=m))
        return self.post_process(x=expansion, device=device)

__init__(name='linear_expansion', c=None, pre_C=None, post_C=None, *args, **kwargs)

The initialization method of the linear expansion function.

It initializes a linear expansion object based on the input function name. This method will also call the initialization method of the base class as well.

Parameters:

Name	Type	Description	Default
name		The name of the linear expansion function.	'linear_expansion'
c		The scalar \(c\) of the linear expansion.	None
pre_C		The \(\mathbf{C}_{pre}\) matrix of the linear expansion.	None
post_C		The \(\mathbf{C}_{post}\) matrix of the linear expansion.	None

Returns:

Type	Description
transformation	The linear expansion function.

Source code in tinybig/expansion/basic_expansion.py

def __init__(self, name='linear_expansion', c=None, pre_C=None, post_C=None, *args, **kwargs):
    r"""
    The initialization method of the linear expansion function.

    It initializes a linear expansion object based on the input function name.
    This method will also call the initialization method of the base class as well.

    Parameters
    ----------
    name: str, default = 'linear_expansion'
        The name of the linear expansion function.
    c: float | torch.Tensor, default = None
        The scalar $c$ of the linear expansion.
    pre_C: torch.Tensor, default = None
        The $\mathbf{C}_{pre}$ matrix of the linear expansion.
    post_C: torch.Tensor, default = None
        The $\mathbf{C}_{post}$ matrix of the linear expansion.

    Returns
    ----------
    transformation
        The linear expansion function.
    """
    super().__init__(name=name, *args, **kwargs)
    self.c = c
    self.pre_C = pre_C
    self.post_C = post_C

calculate_D(m)

The expansion dimension calculation method.

It calculates the intermediate expansion space dimension based on the input dimension parameter m. For the linear expansion function, the expansion space dimension equals to the input space dimension, i.e., $$ D = m. $$

Parameters:

Name	Type	Description	Default
m	int	The dimension of the input space.	required

Returns:

Type	Description
int	The dimension of the expansion space.

Source code in tinybig/expansion/basic_expansion.py

def calculate_D(self, m: int):
    r"""
    The expansion dimension calculation method.

    It calculates the intermediate expansion space dimension based on the input dimension parameter m.
    For the linear expansion function, the expansion space dimension equals to the input space dimension, i.e.,
    $$ D = m. $$

    Parameters
    ----------
    m: int
        The dimension of the input space.

    Returns
    -------
    int
        The dimension of the expansion space.
    """
    return m

forward(x, device='cpu', c=None, pre_C=None, post_C=None, *args, **kwargs)

The forward method of the data expansion function.

It performs the linear data expansion of the input data and returns the expansion result according to one of the following equation: $$ \begin{align} \kappa(\mathbf{x}) &= c \mathbf{x} \in {R}^D, \\ \kappa(\mathbf{x}) &= \mathbf{x} \mathbf{C}_{post} \in {R}^D,\\ \kappa(\mathbf{x}) &= \mathbf{C}_{pre} \mathbf{x} \in {R}^{D}, \end{align} $$ where \(c \in {R}\), \(\mathbf{C}_{post}, \mathbf{C}_{pre} \in {R}^{m \times m}\) denote the provided constant scalar and linear transformation matrices, respectively.

Parameters:

Name	Type	Description	Default
x	Tensor	The input data vector.	required
device		The device to perform the data expansion.	'cpu'
c		The scalar \(c\) of the linear expansion.	None
pre_C		The \(\mathbf{C}_{pre}\) matrix of the linear expansion.	None
post_C		The \(\mathbf{C}_{post}\) matrix of the linear expansion.	None

Returns:

Type	Description
Tensor	The expanded data vector of the input.

Source code in tinybig/expansion/basic_expansion.py

def forward(self, x: torch.Tensor, device='cpu', c=None, pre_C=None, post_C=None, *args, **kwargs):
    r"""
    The forward method of the data expansion function.

    It performs the linear data expansion of the input data and returns the expansion result
    according to one of the following equation:
    $$
    \begin{align}
        \kappa(\mathbf{x}) &= c \mathbf{x} \in {R}^D, \\\\
        \kappa(\mathbf{x}) &= \mathbf{x} \mathbf{C}\_{post} \in {R}^D,\\\\
        \kappa(\mathbf{x}) &= \mathbf{C}\_{pre} \mathbf{x} \in {R}^{D},
    \end{align}
    $$
    where $c \in {R}$, $\mathbf{C}_{post}, \mathbf{C}_{pre} \in {R}^{m \times m}$ denote the provided
    constant scalar and linear transformation matrices, respectively.


    Parameters
    ----------
    x: torch.Tensor
        The input data vector.
    device: str, default = 'cpu'
        The device to perform the data expansion.
    c: float | torch.Tensor, default = None
        The scalar $c$ of the linear expansion.
    pre_C: torch.Tensor, default = None
        The $\mathbf{C}_{pre}$ matrix of the linear expansion.
    post_C: torch.Tensor, default = None
        The $\mathbf{C}_{post}$ matrix of the linear expansion.

    Returns
    ----------
    torch.Tensor
        The expanded data vector of the input.
    """
    b, m = x.shape
    x = self.pre_process(x=x, device=device)

    c = c if c is not None else self.c
    pre_C = pre_C if pre_C is not None else self.pre_C
    post_C = post_C if post_C is not None else self.post_C

    if c is not None:
        expansion = c * x
    elif pre_C is not None:
        assert pre_C.size(-1) == x.size(0)
        expansion = torch.matmul(pre_C, x)
    elif post_C is not None:
        assert x.size(-1) == post_C.size(0)
        expansion = torch.matmul(x, post_C)
    else:
        expansion = x

    assert expansion.shape == (b, self.calculate_D(m=m))
    return self.post_process(x=expansion, device=device)