ó
    "Eñiã	  ã                   ó€   • S SK Jr  S SKJr  S SKJr  S SKJrJr  S/r	\" S5      r
\" S5       " S S\\   5      5       rg	)
é    )ÚSized)ÚTypeVar)Úfunctional_datapipe)Ú	DataChunkÚMapDataPipeÚBatcherMapDataPipeÚ_TÚbatchc                   ó”   ^ • \ rS rSr% Sr\\S'   \\S'   \\S'   S\	4S\\
   S\S\S\\	   SS	4
U 4S
 jjjrS\	4S jrS\4S jrSrU =r$ )r   é   ax  
Create mini-batches of data (functional name: ``batch``).

An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``,
or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``.

Args:
    datapipe: Iterable DataPipe being batched
    batch_size: The size of each batch
    drop_last: Option to drop the last batch if it's not full

Example:
    >>> # xdoctest: +SKIP
    >>> from torchdata.datapipes.map import SequenceWrapper
    >>> dp = SequenceWrapper(range(10))
    >>> batch_dp = dp.batch(batch_size=2)
    >>> list(batch_dp)
    [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]
ÚdatapipeÚ
batch_sizeÚ	drop_lastFÚwrapper_classÚreturnNc                 ót   >• US::  a  [        S5      e[        TU ]	  5         Xl        X l        X0l        X@l        g )Nr   z+Batch size is required to be larger than 0!)ÚAssertionErrorÚsuperÚ__init__r   r   r   r   )Úselfr   r   r   r   Ú	__class__s        €Úd/home/mande/repo/quber/.venv/lib/python3.13/site-packages/torch/utils/data/datapipes/map/grouping.pyr   ÚBatcherMapDataPipe.__init__)   s8   ø€ ð ˜‹?Ü Ð!NÓOÐOÜ‰ÑÔØ ŒØ$ŒØ"ŒØ*Õó    c                 ód  ^ • / n[        UT R                  -  US-   T R                  -  5      n UR                  U 4S jU 5       5        T R                  U5      $ ! [         aK  nT R
                  (       d%  [        U5      S:”  a  T R                  U5      s S nA$ [	        SU S35      UeS nAff = f)Né   c              3   óB   >#   • U  H  nTR                   U   v •  M     g 7f)N)r   )Ú.0Úir   s     €r   Ú	<genexpr>Ú1BatcherMapDataPipe.__getitem__.<locals>.<genexpr><   s   øé € Ð;²7¨a˜Ÿ™ qÖ)²7ùs   ƒr   zIndex z is out of bound.)Úranger   Úextendr   Ú
IndexErrorr   Úlen)r   Úindexr
   ÚindicesÚes   `    r   Ú__getitem__ÚBatcherMapDataPipe.__getitem__8   sž   ø€ ØˆÜ˜ §¡Ñ/°%¸!±)¸t¿¹Ñ1NÓOˆð	KØ�L‰LÔ;±7Ó;Ô;Ø×%Ñ% eÓ,Ð,øÜó 	KØ—>—>¤c¨%£j°1£nØ×)Ñ)¨%Ó0Õ0ä  6¨%¨Ð0AÐ!BÓCÈÐJûð		Kús#   ®+A Á
B/Á$0B*ÂB/ÂB*Â*B/c                 óL  • [        U R                  [        5      (       ae  U R                  (       a"  [	        U R                  5      U R
                  -  $ [	        U R                  5      U R
                  -   S-
  U R
                  -  $ [        [        U 5      R                   S35      e)Nr   z# instance doesn't have valid length)	Ú
isinstancer   r   r   r%   r   Ú	TypeErrorÚtypeÚ__name__)r   s    r   Ú__len__ÚBatcherMapDataPipe.__len__D   sv   € ä�d—m‘m¤U×+Ñ+Ø�~�~Ü˜4Ÿ=™=Ó)¨T¯_©_Ñ<Ð<ä˜DŸM™MÓ*¨T¯_©_Ñ<¸qÑ@ÀTÇ_Á_ÑTÐTäœt D›z×2Ñ2Ð3Ð3VÐWÓXÐXr   )r   r   r   r   )r/   Ú
__module__Ú__qualname__Ú__firstlineno__Ú__doc__r   Ú__annotations__ÚintÚboolr   r	   r.   r   r)   r0   Ú__static_attributes__Ú__classcell__)r   s   @r   r   r      s†   ø‡ ñð( ÓØƒOØƒOð  Ø)2ñ+à˜b‘/ð+ð ð+ð ð	+ð
 ˜I‘ð+ð 
÷+ð +ð
K Iô 
KðY˜÷ Yò Yr   N)Úcollections.abcr   Útypingr   Ú%torch.utils.data.datapipes._decoratorr   Ú#torch.utils.data.datapipes.datapiper   r   Ú__all__r	   r   © r   r   Ú<module>rA      sN   ðå !Ý å Eß Fð  Ð
 €ñ ˆTƒ]€ñ �WÓô<Y˜ YÑ/ó <Yó ñ<Yr   