
    Upf                     &   d Z ddlZddlZddlmZ  ej        dddg          Zej         G d dej                              Z	e	j
        Z
e	j        Ze	j        Z G d	 d
ej                  Zej        Zeee         z  Z G d dej                  ZdS )zPython definitions for third_party.jax.jaxlib.mosaic.python.tpu

These definitions are needed internally by the tpu module.
TODO(tlongeri): Migrate definitions to tpu module
    N)LiteralTargetTuplesublaneslanesc                        e Zd ZdZdZdZd ZdS )	Directionr   r   subelementsc                 4    | j                                         S N)namelowerselfs    `/var/www/html/nettyfy-visnx/env/lib/python3.11/site-packages/jaxlib/mosaic/python/layout_defs.py__repr__zDirection.__repr__    s    9??    N)__name__
__module____qualname__SUBLANESLANESSUBELEMENTSr    r   r   r   r      s2        (
%+    r   r   c                   "    e Zd ZdZd ZeZd ZdS )
Replicated*c                     dS )Nr   r   r   s    r   r   zReplicated.__repr__*   s    3r   c                     dS )NFr   r   s    r   __bool__zReplicated.__bool__.   s    5r   N)r   r   r   
REPLICATEDr   __str__r   r   r   r   r   r   '   s<        *  '    r   r   c                   "    e Zd ZdZdZdefdZdS )ImplicitDimreturnc                 :    t          t          |                     S r   )strintr   s    r   r   zImplicitDim.__repr__8   s    s4yy>>r   N)r   r   r   MINORSECOND_MINORr(   r   r   r   r   r#   r#   4   s9        
%,      r   r#   )__doc__collectionsenumtypingr   
namedtupler   uniqueEnumr   r   r   r   r   r    r)   OffsetIntEnumr#   r   r   r   <module>r5      s   
           $k$]Z4IJJ    	    #        "
	wz"	"    $,     r   