# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/11_batched.ipynb.

# %% auto 0
__all__ = ['BatchedMeta', 'Batched', 'BatchedQMeta', 'BatchedQ']

# %% ../nbs/11_batched.ipynb 4
from typing import Any, Tuple

# %% ../nbs/11_batched.ipynb 5
from .pttype import PTMeta, PTType

# %% ../nbs/11_batched.ipynb 7
class BatchedMeta(PTMeta):
    """
    Metaclass for Batched. It contains logic for instance checking.
    """
    __args__: Tuple[Any]
    _parameterized: bool

    def __getitem__(cls, item):
        return super().__getitem__((item,))
    
    # def __getitem__(cls, item):
    #     if not isinstance(item, int):
    #         raise TypeError("BatchedQ expects an integer argument.")
    #     return super().__getitem__(item)

    # def __instancecheck__(cls, instance):
    #     batch_size = cls.__args__[0]

    #     # Check batch size
    #     if batch_size is not Any and instance.size(0) != batch_size:
    #         return False

    #     return True

# %% ../nbs/11_batched.ipynb 8
class Batched(PTType, metaclass=BatchedMeta):
    """
    A type hint for batched PyTorch tensors.

    Example usage:
    Batched[64]
    """

    __args__ = (Any,)

# %% ../nbs/11_batched.ipynb 10
class BatchedQMeta(PTMeta):
    """
    Metaclass for BatchedQ. It contains logic for instance checking.
    """

    __args__: Tuple[Any]
    _parameterized: bool

    def __getitem__(cls, item):
        return super().__getitem__((item,))
    
    # def __getitem__(cls, item):
    #     if not isinstance(item, int):
    #         raise TypeError("BatchedQ expects an integer argument.")
    #     return super().__getitem__(item)

    # def __instancecheck__(cls, instance):
    #     batch_size = cls.__args__[0]

    #     # Check batch size
    #     if batch_size is not Any and instance.size(0) != batch_size:
    #         return False

    #     return True

# %% ../nbs/11_batched.ipynb 11
class BatchedQ(PTType, metaclass=BatchedQMeta):
    pass
    """
    A type hint for batched PyTorch tensors with an unknown batch size.

    Example usage:
    BatchedQ[64]
    """

    __args__ = (Any,)
