_matrix.py

from typing import TypeGuard

from torch import Tensor


class GeneralizedMatrix(Tensor):
    pass


class Matrix(GeneralizedMatrix):
    pass


class PSDGeneralizedMatrix(Tensor):
    pass


class PSDMatrix(PSDGeneralizedMatrix, Matrix):
    pass


def is_generalized_matrix(t: Tensor) -> TypeGuard[GeneralizedMatrix]:
    return t.ndim >= 1


def is_matrix(t: Tensor) -> TypeGuard[Matrix]:
    return t.ndim == 2


def is_psd_generalized_matrix(t: Tensor) -> TypeGuard[PSDGeneralizedMatrix]:
    half_dim = t.ndim // 2
    return t.ndim % 2 == 0 and t.shape[:half_dim] == t.shape[: half_dim - 1 : -1]
    # We do not check that t is PSD as it is expensive, but this must be checked in the tests of
    # every function that use this TypeGuard by using `assert_psd_generalized_matrix`.


def is_psd_matrix(t: Tensor) -> TypeGuard[PSDMatrix]:
    return t.ndim == 2 and t.shape[0] == t.shape[1]
    # We do not check that t is PSD as it is expensive, but this must be checked in the tests of
    # every function that use this TypeGuard, by using `assert_psd_matrix`.