
    `iV                         d dl mZmZ d dlmZmZ d dlmZmZ d dl	m
Z
  G d d          Z G d de          Zi Zd	 Zd
 ZefdZd ZdS )   )reifyunify)firstgroupby)	_toposortfreezeisvarc                   ,    e Zd Zd Zd Zd Zd Zd ZdS )
Dispatcherc                 0    || _         i | _        g | _        d S N)namefuncsordering)selfr   s     {/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/torch/fx/experimental/unification/match.py__init__zDispatcher.__init__	   s    	
    c                 f    || j         t          |          <   t          | j                   | _        d S r   )r   r   r   )r   	signaturefuncs      r   addzDispatcher.add   s+    (,
6)$$% ,,r   c                 B    |                      |          \  }} ||i |S r   )resolve)r   argskwargsr   _s        r   __call__zDispatcher.__call__   s,    ,,t$$atT$V$$$r   c                 2   t          |          }| j        D ]J}t          |          |k    rt          t          |          |          }|dur| j        |         }||fc S Kt          dt          | j                  z   dz   t          |          z             )NFz No match found. 
Known matches: z
Input: )lenr   r   r   r   NotImplementedErrorstr)r   r   nr   sresults         r   r   zDispatcher.resolve   s    II 	! 	!I9~~""fTllI..A~~I.qy     "/$-  ! $ii
 
 	
r   c                       fd}|S )Nc                 4                         |            S r   )r   )r   r   r   s    r   r   zDispatcher.register.<locals>._'   s    HHY%%%Kr    )r   r   r   s   `` r   registerzDispatcher.register&   s)    	 	 	 	 	 	 r   N)__name__
__module____qualname__r   r   r   r   r*   r)   r   r   r   r      s_          
- - -% % %
 
 
     r   r   c                       e Zd ZdZd ZdS )VarDispatcheraX  A dispatcher that calls functions with variable names
    >>> # xdoctest: +SKIP
    >>> d = VarDispatcher("d")
    >>> x = var("x")
    >>> @d.register("inc", x)
    ... def f(x):
    ...     return x + 1
    >>> @d.register("double", x)
    ... def f(x):
    ...     return x * 2
    >>> d("inc", 10)
    11
    >>> d("double", 10)
    20
    c                 ~    |                      |          \  }}d |                                D             } |di |S )Nc                 $    i | ]\  }}|j         |S r)   )token.0kvs      r   
<dictcomp>z*VarDispatcher.__call__.<locals>.<dictcomp>A   s     ...DAqQWa...r   r)   )r   items)r   r   r   r   r%   ds         r   r   zVarDispatcher.__call__?   sD    ,,t$$a..AGGII...tyyayyr   N)r+   r,   r-   __doc__r   r)   r   r   r/   r/   .   s-              r   r/   c                       |                     dt                    |                     dt                     fd}|S )N	namespacer   c                 v    | j         }|vr |          |<   |         }|                    |            |S r   )r+   r   )r   r   r9   
dispatcherr<   r   s      r   r   zmatch.<locals>._L   sJ    }y  (j..IdOdO	ir   )getglobal_namespacer   )r   r   r   r>   r<   s   `  @@r   matchrA   H   sV    

;(899IL*55J	 	 	 	 	 	 	 Hr   c                    t          |          rt          |           sdS t          | |          }|du rdS d |                                D             }t          | |          | k    rdS t          ||          |k    rdS dS )z)``a`` is a more specific match than ``b``TFc                 V    i | ]&\  }}t          |          rt          |          #||'S r)   r	   r3   s      r   r7   zsupercedes.<locals>.<dictcomp>a   s4    DDD$!Q588D588DADDDr   N)r
   r   r8   r   )abr%   s      r   
supercedesrF   Z   s    Qxx a taAEzzuDD!''))DDDAQ{{atQ{{au r   c                 z    t          | |          r*t          ||           r ||            ||          k    S dS dS )zUA should be checked before B
    Tie broken by tie_breaker, defaults to ``hash``
    TF)rF   )rD   rE   tie_breakers      r   edgerI   i   sM     !Q a 	;q>>KKNN2245r   c                     t          t          t                                 fd D             }t          t          |          } D ]}||vrg ||<   d |                                D             }t          |          S )zA sane ordering of signatures to check, first to last
    Topological sort of edges as given by ``edge`` and ``supercedes``
    c                 B    g | ]}D ]}t          ||          ||fS r)   )rI   )r4   rD   rE   
signaturess      r   
<listcomp>zordering.<locals>.<listcomp>{   s7    JJJJJAtAqzzJaVJJJJr   c                 .    i | ]\  }}|d  |D             S )c                     g | ]\  }}|S r)   r)   )r4   rD   rE   s      r   rM   z'ordering.<locals>.<dictcomp>.<listcomp>   s    !!!tq!!!!r   r)   r3   s      r   r7   zordering.<locals>.<dictcomp>   s-    <<<daQ!!q!!!<<<r   )listmaptupler   r   r8   r   )rL   edgesr%   s   `  r   r   r   v   s     c%,,--JJJJJZJJJEE5!!E  E>>E!H<<ekkmm<<<EUr   N)corer   r   unification_toolsr   r   utilsr   r   variabler
   r   r/   r@   rA   rF   hashrI   r   r)   r   r   <module>rY      s
           - - - - - - - - $ $ $ $ $ $ $ $      # # # # # # # #L    J   .    $     	 	 	 	    r   