
    sg"                    t   U d dl mZ d dlZd dlZd dlmZmZ d dlZd dlm	Z	 d dl
Z
d dlZ
d dlmZ d dlmZ d dlmZmZ d dlmZmZ erd dlZd	Zd
ed<   dZe
j6                  j9                  de      Zd,dZej>                  d-d       Z d-dZ!e jD                  d.d       Z#e jD                  d/d       Z$e jD                  d0d       Z%e jD                  d1d       Z&e jD                  d2d       Z'e jD                  d3d       Z(e jD                  d4d       Z)e jD                  d5d       Z*e jD                  d6d       Z+e jD                  d7d       Z,e jD                  d8d       Z-e jD                  d9d       Z.e jD                  d:d       Z/e jD                  d;d       Z0e jD                  d<d       Z1e jD                  d=d        Z2e jD                  d>d!       Z3e jD                  d?d"       Z4e jD                  d@d#       Z5dAd$Z6ejn                  Z7ejp                  Z8ejr                  Z9ejt                  Z:ejv                  Z;ejx                   G d% d&ejz                               Z=ejx                   G d' d(ej|                  e=                Z> ej~                  ej                  e=e!)      Z@ejx                   G d* d+e=             ZAy)B    )annotationsN)AnyTYPE_CHECKING)graph_building)diagnostics)infra)	decorator	formatter)registration
type_utils
   int_CONTAINER_ITEM_LIMITonnx_diagnosticsz
torch.onnxc                 p    t         j                  j                  j                  j	                  t
              S N)torch_logging	_internal	log_stateis_artifact_enabled&_ONNX_DIAGNOSTICS_ARTIFACT_LOGGER_NAME     V/var/www/html/venv/lib/python3.12/site-packages/torch/onnx/_internal/fx/diagnostics.py(is_onnx_diagnostics_log_artifact_enabledr   .   s(    >>##--AA. r   c                ,    t        j                  |       S r   )r
   format_argumentobjs    r   _format_argumentr!   4   s    $$S))r   c                N    t         j                  t        |             } ||       S r   )r!   dispatchtype)r    r
   s     r   r   r   9   s      ))$s)4IS>r   c                6    d| j                   j                   dS )Nztorch.nn.Module()	__class____name__r   s    r   _torch_nn_moduler*   G   s    cmm445Q77r   c                6    d| j                   j                   dS )Nztorch.fx.GraphModule(r&   r'   r   s    r   _torch_fx_graph_moduler,   L   s    "3==#9#9":!<<r   c                    d| j                    d| j                   d}d| j                  vr|dz   S |t        | j                  d         z   S )Nzfx.Node(z)[z]:valNone)targetopmetar   )r    node_strings     r   _torch_fx_noder4   Q   sM    SZZL366("5KCHHV##%999r   c                    d|  dS )NzSymBool(r&   r   r   s    r   _torch_fx_symbolic_boolr6   Y   s    cU!r   c                    d|  dS )NzSymInt(r&   r   r   s    r   _torch_fx_symbolic_intr8   ^   s    SEr   c                    d|  dS )Nz	SymFloat(r&   r   r   s    r   _torch_fx_symbolic_floatr:   c   s    se1r   c                r    dt        j                  | j                         t        | j                         dS )NzTensor(r&   fx_type_utilsfrom_torch_dtype_to_abbrdtype_stringify_shapeshaper   s    r   _torch_tensorrB   h   s4    ];;CIIFGHXY\YbYbHcGddeffr   c                    t        |       S r   strr   s    r   _intrF   m       s8Or   c                    t        |       S r   rD   r   s    r   _floatrI   r   rG   r   c                    t        |       S r   rD   r   s    r   _boolrK   w   rG   r   c                    | S r   r   r   s    r   _strrM   |   s    Jr   c                V    d| j                    d| j                   d| j                   dS )Nzregistration.ONNXFunction(z, is_custom=z, is_complex=r&   )op_full_name	is_custom
is_complexr   s    r   _registration_onnx_functionrR      s4     ((8(8'9cmm_Tabebpbpaqqrssr   c                    dt        |        d}| s|dz   S t        |       D ]*  \  }}|t        k\  r|dz  } |dz   S |t        |       dz  }, |dz   S )NzList[length=](
None)...,
,
r&   len	enumerater   r   )r    list_stringiitems       r   _listr^      s     S
$/KW$$S> 54%%8#K 	/$/0445 r   c                    dt        |        d}| s|dz   S t        |       D ]*  \  }}|t        k\  r|dz  } |dz   S |t        |       dz  }, |dz   S )NzTuple[length=rT   rU   rV   rW   r&   rX   )r    tuple_stringr\   r]   s       r   _tuplera      s    "3s8*D1Lg%%S> 64%%H$L# 	?401556 #r   c                    dt        |        d}| s|dz   S t        | j                               D ]0  \  }\  }}|t        k\  r|dz  } |dz   S || dt	        |       dz  }2 |dz   S )NzDict[length=rT   rU   z...
z: rW   r&   )rY   rZ   itemsr   r   )r    dict_stringr\   keyvalues        r   _dictrg      s     S
$/KW$$$SYY[1 =<C%%7"K 	#b!7 8<<= r   c                4    dt        | j                         dS )Nz
Parameter(r&   )r   datar   s    r   _torch_nn_parameterrj      s    12!44r   c                r    dt        j                  | j                         t        | j                         dS )Nz`TorchScriptTensor()`r<   r   s    r   _onnxscript_torch_script_tensorrm      s4     !G!G		!R STdehenenToSpprssr   c                "    d| j                    dS )Nz`OnnxFunction(rl   namer   s    r   _onnxscript_onnx_functionrq      s    CHH:R((r   c                "    d| j                    dS )Nz`TracedOnnxFunction(rl   ro   r   s    r    _onnxscript_traced_onnx_functionrs      s    !#((2..r   c                @    | yddj                  d | D               dS )N [z, c              3  2   K   | ]  }t        |        y wr   rD   ).0xs     r   	<genexpr>z#_stringify_shape.<locals>.<genexpr>   s     /AQ/s   ])join)rA   s    r   r@   r@      s)    }tyy///022r   c                  H    e Zd ZU  ej                  de      Zded<   ddZy)
DiagnosticFinitdefaultlogging.Loggerloggerc                    | j                   j                  |      rI||z  }t               r | j                   j                  ||fi | | j                  j                  |       y y r   )r   isEnabledForr   logadditional_messagesappend)selflevelmessageargskwargsformatted_messages         r   r   zDiagnostic.log   sZ    ;;##E* '$79  '8CFC$$++,=> +r   N)r   r   r   rE   returnr/   )	r)   
__module____qualname__dataclassesfielddiagnostic_loggerr   __annotations__r   r   r   r   r~   r~      s!    .[..ECTUFNU?r   r~   c                       e Zd ZU  ej                  de      Zded<    ej                  de      Z	ded<    fdZ
 xZS )DiagnosticContextFr   r   r   ztype[Diagnostic]_bound_diagnostic_typec                n    | j                   j                  | _        t               st        |          S | S r   )r   r   _previous_log_levelr   super	__enter__r   r(   s    r   r   zDiagnosticContext.__enter__   s0    #';;#4#4  897$&&Kr   )r)   r   r   r   r   r   r   r   r~   r   r   __classcell__r(   s   @r   r   r      sG    .[..ECTUFNU/@{/@/@J0,  r   r   )diagnostic_typer   c                  0     e Zd ZU dZded<   d fdZ xZS )UnsupportedFxNodeDiagnosticNztorch.fx.Node | Noneunsupported_fx_nodec                P    t         |           | j                  t        d      y )Nz&unsupported_fx_node must be specified.)r   __post_init__r   
ValueErrorr   s    r   r   z)UnsupportedFxNodeDiagnostic.__post_init__   s-    
 ##+EFF ,r   )r   r/   )r)   r   r   r   r   r   r   r   s   @r   r   r      s    04-4G Gr   r   )r   bool)r    r   r   rE   )r    ztorch.nn.Moduler   rE   )r    ztorch.fx.GraphModuler   rE   )r    ztorch.fx.Noder   rE   )r    ztorch.SymBoolr   rE   )r    ztorch.SymIntr   rE   )r    ztorch.SymFloatr   rE   )r    ztorch.Tensorr   rE   )r    r   r   rE   )r    floatr   rE   )r    r   r   rE   )r    rE   r   rE   )r    zregistration.ONNXFunctionr   rE   )r    listr   rE   )r    tupler   rE   )r    dictr   rE   )r    ztorch.nn.Parameterr   rE   )r    z graph_building.TorchScriptTensorr   rE   )r    zonnxscript.OnnxFunctionr   rE   )r    zonnxscript.TracedOnnxFunctionr   rE   )rA   ztorch.Size | Noner   rE   )B
__future__r   r   	functoolstypingr   r   
onnxscript"onnxscript.function_libs.torch_libr   r   torch.fxtorch.onnx._internalr    torch.onnx._internal.diagnosticsr   &torch.onnx._internal.diagnostics.infrar	   r
   torch.onnx._internal.fxr   r   r=   loggingr   r   r   r   getArtifactLoggerr   r   singledispatchr!   r   registerr*   r,   r4   r6   r8   r:   rB   rF   rI   rK   rM   rR   r^   ra   rg   rj   rm   rq   rs   r@   ruleslevelsRuntimeErrorWithDiagnostic
LazyStringDiagnosticOptions	dataclassr~   r   partialdiagnose_callr   r   r   r   <module>r      sS   "   %  =   , 2 G M    s  *< &NN448 
 * * 8 8 = = : :       g g         t t
 
 
 
 
 
 
 5 5 t t ) ) / /
3 				"== !!
++  ?!! ? ? //
;    "	!!# 
G* 
G 
Gr   