U
    9%e<                      @   s4  d dl Z d dlmZ d dlmZmZ d dlZd dlm	Z	m
Z
mZmZmZmZmZmZmZmZmZmZ d dlmZmZ d dlmZmZmZ d dlmZ d dlmZm Z m!Z! d d	l"m#Z#m$Z$m%Z%m&Z& d d
l'm(Z( d dl)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5 d dl6m7Z7m8Z8m9Z9m:Z: d dl;m<Z<m=Z=m>Z>m?Z?m@Z@mAZAmBZBmCZCmDZDmEZEmFZFmGZG d dlHmIZI dddddddddddddddddd d!d"d#d$d%d&d'gZJeZKeZLe&ZMd(d ZNejOjPZQeIeQd) deQ_RdeQ_Sd*d ZTd+d! ZUd,d ZVG d-d# d#ZWeed d.d/d0ZXeYd1d2dZZd3d Z[[ejO\ s0e]d4dS )5    N)contextmanager)AnyIterator)_Await_drop_IgnoreContextManager_isinstance	_overload_overload_methodexportFinalFutureignoreis_scriptingunused)forkwait)
_awaitable_awaitable_nowait_awaitable_wait)_register_decomposition)freezeoptimize_for_inferencerun_frozen_optimizations)fuserlast_executed_optimized_graphoptimized_executionset_fusion_strategy)_InsertPoint)_ScriptProfile_unwrap_optional	AttributeCompilationUnit	interfaceRecursiveScriptClassRecursiveScriptModulescriptscript_methodScriptFunctionScriptModuleScriptWarning)jit_module_from_flatbufferloadsavesave_jit_module_to_flatbuffer)_flatten_get_trace_graph_script_if_tracing_unique_state_dict
is_tracingONNXTracedModuleTopLevelTracedModuletracetrace_moduleTracedModuleTracerWarningTracingCheckError)
set_moduler!   r"   Errorr   r(   r)   annotateenable_onednn_fusionexport_opnamesr   r   r   
isinstancer,   onednn_fusion_enabledr   r-   r&   script_if_tracingr   strict_fusionr6   r7   r   r   c                 C   s   t j| jS )a  
    Generates new bytecode for a Script module and returns what the op list
    would be for a Script Module based off the current code base. If you
    have a LiteScriptModule and want to get the currently present
    list of ops call _export_operator_list instead.
    )torch_CZ_export_opnamesZ_c)m rG   Q/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/torch/jit/__init__.pyr?   i   s    z	torch.jitc                 C   s   |S )a6  
    This method is a pass-through function that returns `the_value`, used to hint TorchScript
    compiler the type of `the_value`. It is a no-op when running outside of TorchScript.

    Though TorchScript can infer correct type for most Python expressions, there are some cases where
    type inference can be wrong, including:

    - Empty containers like `[]` and `{}`, which TorchScript assumes to be container of `Tensor`
    - Optional types like `Optional[T]` but assigned a valid value of type `T`, TorchScript would assume
      it is type `T` rather than `Optional[T]`

    Note that `annotate()` does not help in `__init__` method of `torch.nn.Module` subclasses because it
    is executed in eager mode. To annotate types of `torch.nn.Module` attributes,
    use :meth:`~torch.jit.Annotate` instead.

    Example:

    .. testcode::

        import torch
        from typing import Dict

        @torch.jit.script
        def fn():
            # Telling TorchScript that this empty dictionary is a (str -> int) dictionary
            # instead of default dictionary type of (str -> Tensor).
            d = torch.jit.annotate(Dict[str, int], {})

            # Without `torch.jit.annotate` above, following statement would fail because of
            # type mismatch.
            d["name"] = 20

    .. testcleanup::

        del fn

    Args:
        the_type: Python type that should be passed to TorchScript compiler as type hint for `the_value`
        the_value: Value or expression to hint type for.

    Returns:
        `the_value` is passed back as return value.
    rG   )Zthe_typeZ	the_valuerG   rG   rH   r=   |   s    ,c                 C   s   t | S )a  
    Compiles ``fn`` when it is first called during tracing. ``torch.jit.script``
    has a non-negligible start up time when it is first called due to
    lazy-initializations of many compiler builtins. Therefore you should not use
    it in library code. However, you may want to have parts of your library work
    in tracing even if they use control flow. In these cases, you should use
    ``@torch.jit.script_if_tracing`` to substitute for
    ``torch.jit.script``.

    Args:
        fn: A function to compile.

    Returns:
        If called during tracing, a :class:`ScriptFunction` created by `torch.jit.script` is returned.
        Otherwise, the original function `fn` is returned.
    )r1   )fnrG   rG   rH   rB      s    c                 C   s
   t | |S )ai  
    This function provides for container type refinement in TorchScript. It can refine
    parameterized containers of the List, Dict, Tuple, and Optional types. E.g. ``List[str]``,
    ``Dict[str, List[torch.Tensor]]``, ``Optional[Tuple[int,str,int]]``. It can also
    refine basic types such as bools and ints that are available in TorchScript.

    Args:
        obj: object to refine the type of
        target_type: type to try to refine obj to
    Returns:
        ``bool``: True if obj was successfully refined to the type of target_type,
            False otherwise with no new type refinement


    Example (using ``torch.jit.isinstance`` for type refinement):
    .. testcode::

        import torch
        from typing import Any, Dict, List

        class MyModule(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, input: Any): # note the Any type
                if torch.jit.isinstance(input, List[torch.Tensor]):
                    for t in input:
                        y = t.clamp(0, 0.5)
                elif torch.jit.isinstance(input, Dict[str, str]):
                    for val in input.values():
                        print(val)

        m = torch.jit.script(MyModule())
        x = [torch.rand(3,3), torch.rand(4,3)]
        m(x)
        y = {"key1":"val1","key2":"val2"}
        m(y)
    )r   )objZtarget_typerG   rG   rH   r@      s    'c                   @   s4   e Zd ZdZdd Zdd Zeeedddd	ZdS )
rC   aC  
    This class errors if not all nodes have been fused in
    inference, or symbolically differentiated in training.

    Example:

    Forcing fusion of additions.

    .. code-block:: python

        @torch.jit.script
        def foo(x):
            with torch.jit.strict_fusion():
                return x + x + x

    c                 C   s   t j std d S )NzOnly works in script mode)rD   Z_jit_internalr   warningswarnselfrG   rG   rH   __init__   s    

zstrict_fusion.__init__c                 C   s   d S NrG   rM   rG   rG   rH   	__enter__  s    zstrict_fusion.__enter__N)typevaluetbreturnc                 C   s   d S rP   rG   )rN   rR   rS   rT   rG   rG   rH   __exit__  s    zstrict_fusion.__exit__)__name__
__module____qualname____doc__rO   rQ   r   rV   rG   rG   rG   rH   rC      s   )rU   c               	   c   s8   t jjj} zt jjd d V  W 5 t jj|  X d S )NF)rD   rE   ZGraphZglobal_print_source_rangesZset_global_print_source_ranges)Zold_enable_source_rangesrG   rG   rH   _hide_source_ranges  s
    

r[   enabledc                 C   s   t j|  dS )zQ
    Enables or disables onednn JIT fusion based on the parameter `enabled`.
    N)rD   rE   Z_jit_set_llga_enabledr\   rG   rG   rH   r>     s    c                   C   s
   t j S )z6
    Returns whether onednn JIT fusion is enabled
    )rD   rE   Z_jit_llga_enabledrG   rG   rG   rH   rA     s    zJIT initialization failed)^rK   
contextlibr   typingr   r   Ztorch._CrD   Ztorch._jit_internalr   r   r   r   r	   r
   r   r   r   r   r   r   Ztorch.jit._asyncr   r   Ztorch.jit._awaitr   r   r   Ztorch.jit._decomposition_utilsr   Ztorch.jit._freezer   r   r   Ztorch.jit._fuserr   r   r   r   Ztorch.jit._ir_utilsr   Ztorch.jit._scriptr   r    r!   r"   r#   r$   r%   r&   r'   r(   r)   r*   Ztorch.jit._serializationr+   r,   r-   r.   Ztorch.jit._tracer/   r0   r1   r2   r3   r4   r5   r6   r7   r8   r9   r:   Ztorch.utilsr;   __all__Z_fork_waitZ_set_fusion_strategyr?   rE   ZJITExceptionr<   rW   rY   r=   rB   r@   rC   r[   boolr>   rA   Z	_jit_initRuntimeErrorrG   rG   rG   rH   <module>   sv   888
/*!	