afnio.cognitive.modules.split

class Split(Module):
    """
    Splits a single input Variable into multiple output Variables.

    This module utilizes the [`Split`][afnio.autodiff.basic_ops.Split] operation
    from `afnio.autodiff.basic_ops`. It supports string data types, splitting the string
    data of the input Variable based on a specified delimiter and an optional maximum
    number of splits.

    Note:
        This module does not have any trainable parameters.

    Examples:
        >>> from afnio import cognitive as cog
        >>> class Splitter(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.split = cog.Split()
        >>>     def forward(self, x):
        ...         return self.split(x, " ", 1)
        >>> input = afnio.Variable(data="Afnio is great!", role="sentence")
        >>> splitter = Splitter()
        >>> result = splitter(input)
        >>> print([r.data for r in result])
        ['Afnio', 'is great!']
        >>> print([r.role for r in result])
        ['split part 0 of sentence', 'split part 1 of sentence']

    Raises:
        TypeError: If `x` is not an instance of [`Variable`][afnio.Variable] which
            [`data`][afnio.Variable.data] attribute is a string or a list of strings,
            or if `sep` is not a string or `Variable` containing a string,
            or if `maxsplit` is not an integer or `Variable` containing an integer.

    See Also:
        [`afnio.autodiff.basic_ops.Split`][afnio.autodiff.basic_ops.Split]
        for the underlying operation.
    """

    sep: Optional[Union[str, Variable]]
    maxsplit: Optional[Union[int, Variable]]

    def __init__(self):
        super().__init__()

        self.register_buffer("sep", None)
        self.register_buffer("maxsplit", None)

    def forward(
        self,
        x: Variable,
        sep: Optional[Union[str, Variable]] = None,
        maxsplit: Optional[Union[int, Variable]] = -1,
    ) -> List[Variable]:
        """
        Forward pass for splitting a `Variable`.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            x: The input `Variable` to be split.
            sep: The delimiter to use for splitting the string. If `None`, splits on
                whitespace. Can be a string or a `Variable` containing a string.
            maxsplit: The maximum number of splits to perform. If `-1`, there is no
                limit on the number of splits. Can be an integer or a `Variable`
                containing an integer.

        Returns:
            A tuple of `Variable` instances resulting from the split operation, \
            each with appropriately assigned [`data`][afnio.Variable.data], \
            [`role`][afnio.Variable.role], and \
            [`requires_grad`][afnio.Variable.requires_grad] attributes.

        Raises:
            TypeError: If `x` is not an instance of [`Variable`][afnio.Variable] which
                [`data`][afnio.Variable.data] attribute is a string or a list of
                strings, or if `sep` is not a string or `Variable` containing a string,
                or if `maxsplit` is not an integer or `Variable` containing an integer.
        """
        self.sep = (
            None
            if sep is None
            else (sep if isinstance(sep, Variable) else Variable(sep))
        )
        self.maxsplit = (
            None
            if maxsplit is None
            else (maxsplit if isinstance(maxsplit, Variable) else Variable(maxsplit))
        )
        return SplitOp.apply(x, self.sep, self.maxsplit)