o
    "i                     @   s   U d dl mZ ddlmZ ddlZddlZddlZg dZe Z	ee
 ed< dgZdd	 ZG d
d dZdd ZdddZdddZdS )   )
OpOverload    )SetN)Libraryimpldefinefallthrough_kernel_implsZprimc                   C   s   t d)zZ
    A dummy function to pass to ``Library.impl`` in order to register a fallthrough.
    z,fallthrough_kernel() should never be called.)NotImplementedError r   r   \/var/www/html/eduruby.in/lip-sync/lip-sync-env/lib/python3.10/site-packages/torch/library.pyr      s   r   c                   @   s6   e Zd ZdZdddZdd ZdddZdd	d
ZdS )r   aC  
    A class to create libraries that can be used to register new operators or
    override operators in existing libraries from Python.
    A user can optionally pass in a dispatch keyname if they only want to register
    kernels corresponding to only one specific dispatch key.

    To create a library to override operators in an existing library (with name ns), set the kind to "IMPL".
    To create a new library (with name ns) to register new operators, set the kind to "DEF".
    To create a fragment of a possibly existing library to register operators (and bypass
    the limitation that there is only one library for a given namespace), set the kind to
    "FRAGMENT".

    Args:
        ns: library name
        kind: "DEF", "IMPL" (default: "IMPL"), "FRAGMENT"
        dispatch_key: PyTorch dispatch key (default: "")
     c                 C   s   |dvr	t d||tv r|dks|dkrt |dtjddd }|j|j}}tj|||||| _	|| _
t | _|| _|| _t| tt| j d S )	N)ZIMPLDEFFRAGMENTzUnsupported kind: r   r   zJ is a reserved namespace. Please try creating a library with another name.   )limitr   )
ValueError_reserved_namespaces	tracebackextract_stackfilenamelinenotorch_CZ_dispatch_librarymnsset	_op_implskinddispatch_keyweakreffinalize_del_libraryr	   )selfr   r   r   framer   r   r   r   r   __init__*   s   

zLibrary.__init__c                 C   s   d| j  d| j d| j dS )NzLibrary(kind=z, ns=z, dispatch_key=z)>)r   r   r   )r#   r   r   r   __repr__>   s   zLibrary.__repr__c                 C   s$   |dvrt d| | j||S )az  Defines a new operator and its semantics in the ns namespace.

        Args:
            schema: function schema to define a new operator.
            alias_analysis (optional): Indicates if the aliasing properties of the operator arguments can be
                                       inferred from the schema (default behavior) or not ("CONSERVATIVE").
        Returns:
            name of the operator as inferred from the schema.

        Example::
            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LIBRARY)
            >>> my_lib = Library("foo", "DEF")
            >>> my_lib.define("sum(Tensor self) -> Tensor")
        )r   ZFROM_SCHEMAZCONSERVATIVEzInvalid alias_analysis type )RuntimeErrorr   r   )r#   schemaalias_analysisr   r   r   r   A   s   zLibrary.definec                 C   s.  t |stdt| |dkr| j}t|tr|}nt|tr4|jj}|jj	}|dkr3|d | }nt
d| jd |dd  d | }|tv r[t
d|dd || j|d	kr||}d|vrm| j d| }tj|d
r|t
d| d| j||dkr|nd
| t| | j| dS )a  Registers the function implementation for an operator defined in the library.

        Args:
            op_name: operator name (along with the overload) or OpOverload object.
            fn: function that's the operator implementation for the input dispatch key or :func:`~fallthrough_kernel`
                to register a fallthrough.
            dispatch_key: dispatch key that the input function should be registered for. By default, it uses
                          the dispatch key that the library was created with.

        Example::
            >>> my_lib = Library("aten", "IMPL")
            >>> def div_cpu(self, other):
            >>>     return self * (1 / other)
            >>> my_lib.impl("div.Tensor", div_cpu, "CPU")
        z;Input function is required to be a callable but found type r   .zQimpl should be passed either a name or an OpOverload object as the first argument/z::zThis is not allowed since there's already a kernel registered from python overriding {}'s behavior for {} dispatch key and {} namespace.ZMetaZCompositeImplicitAutogradz?We should not register a meta kernel directly to the operator 'z', because it has a CompositeImplicitAutograd kernel in core. Instead we should let the operator decompose, and ensure that we have meta kernels for the base ops that it decomposes into.N)callable	TypeErrortyper   
isinstancestrr   Z_schemanameoverload_namer'   r   splitr	   formatr   r   Z%_dispatch_has_kernel_for_dispatch_keyr   r   addr   )r#   Zop_namefnr   r2   r3   keyZdispatcher_op_namer   r   r   r   V   s:   

 

zLibrary.implNr   )__name__
__module____qualname____doc__r%   r&   r   r   r   r   r   r   r      s    

r   c                 C   s   | |8 } d S Nr   )Zcaptured_implsZop_implsr   r   r   r"      s   r"   r   c                        fdd}|S )Nc                    s    |   | S r>   )r   )fr   libr2   r   r   wrap   s   zimpl.<locals>.wrapr   )rB   r2   r   rC   r   rA   r   r      s   r   c                    r?   )Nc                    s     }||  | S r>   )r   r   )r@   r2   r)   rB   r(   r   r   rC      s   zdefine.<locals>.wrapr   )rB   r(   r)   rC   r   rD   r   r      s   r   r9   )Z_opsr   typingr   r   r   r    __all__r   r	   r1   __annotations__r   r   r   r"   r   r   r   r   r   r   <module>   s    {
