
    Upf\                     $   d Z ddlT ddlmZ ddlT 	 ddlmZ n# e$ r	 ddlmZ Y nw xY wej        	                    d            ej
        ed           G d	 d
e                      Z ej
        ed           G d de                      ZdS )z)Python bindings for the MLIR TPU dialect.   )*)_Dialect    )_cextzjax.jaxlib.mosaic.pythonT)replacec                   @     e Zd ZdZddd fd
Zed             Z xZS )TraceOpz=An extension to the automatically generated TraceOp bindings.Nlocipc                    t                                          |||||            | j        d         j        j        g   d S Nr
   r   super__init__regionsblocksappend)selfresultsmessagelevelr   r   	__class__s         X/var/www/html/nettyfy-visnx/env/lib/python3.11/site-packages/jaxlib/mosaic/python/tpu.pyr   zTraceOp.__init__&   sE    	GGWgu#"===!DLO!2&&&&    c                 2    | j         d         j        d         S Nr   r   r   r   s    r   bodyzTraceOp.body*       <?!!$$r   __name__
__module____qualname____doc__r   propertyr    __classcell__r   s   @r   r	   r	   "   sh        EE59d ' ' ' ' ' ' ' % % 8% % % % %r   r	   c                   @     e Zd ZdZddd fd
Zed             Z xZS )RegionOpz>An extension to the automatically generated RegionOp bindings.Nr
   c                    t                                          |||           | j        d         j                                         d S r   r   )r   r   r   r   r   s       r   r   zRegionOp.__init__3   sC    	GGW#"---LO!!#####r   c                 2    | j         d         j        d         S r   r   r   s    r   r    zRegionOp.body7   r!   r   r"   r)   s   @r   r+   r+   /   sh        FF%)d $ $ $ $ $ $ $ % % 8% % % % %r   r+   N)r&   _tpu_genr   jaxlib.mlir._mlir_libs._tpu_ext jaxlib.mlir.dialects._ods_commonr   ImportErrormlir.dialects._ods_commonglobalsappend_dialect_search_prefixregister_operationr	   r+    r   r   <module>r7      s;   0 /           - - - -.4444444 . . .--------.  * *+E F F F (D111	% 	% 	% 	% 	%g 	% 	% 21	% (D111	% 	% 	% 	% 	%x 	% 	% 21	% 	% 	%s    ''