a OSicc@sddlmZddlmZmZmZmZmZmZm Z m Z ddl m Z ddl mZmZddlZddlZeejdrddlmZejejejejejejejejejejej ej!ej"ej#ej$ej%ej&ej'i Z(niZ(dd d d Z)eej*e e$e e$d ffZ+ee e$e e$d ffZ,ee$e e$e e$d ffZ-ee e$e e$d ffZ.ee&e$ee/fZ0e&e$ee/fZ1Gd ddej2Z3eej2e3fZ4ej2e3fZ5ee e4e e4d ffZ6dddddZ7ddddddZ8dddddZ9ddddZ:dd d!d"Z;d#d$d%d&Zdd.d/d0d1Z?ddd2d3d4Z@dddd5d6d7ZAdd8d8d9d:d;ZBdd.dd<d=d>ZCd?d?ddd@dAZDdBdddCdDZEdEdFZFdGdHZGejHejIejJejKejLfZMejNej ejOejPfZQejRejSejTfZUdddIdJdKZVdddIdLdMZWdddIdNdOZXdddIdPdQZYejTejPejSejOejRejNiZZejNejRej ejSejOejSejPejTiZ[dddIdRdSZ\dddIdTdUZ]ddVdIdWdXZ^e&ej&e$ejLeejPe/ejTiZ_dVddYdZd[Z`e&e$ee/fZadVdVdVdd\d]Zbd^d^d_dd`daZcdVdVdddbdcZdddddddedfZedgdhZfej ejOejNejOejRejSiZgdddIdidjZhGdkdldle ZidldmdndodpZjdqdrdsdtduZkd#ddvdwdxZld#d?ddydzd{Zmd?d|d}d~Znd#dddddZodd#d'ddddZpdS)) annotations)AnyUnionSequenceOptionalCallableDictTupleList)Enum)reduce cmp_to_keyN_nvfuser)DataTypez torch.dtypedtypecCst|S)z@ Translates from torch.dtype to nvFuser's DataType enum )!_torch_dtype_to_nvfuser_dtype_maprrN/var/www/html/django/DPS/env/lib/python3.9/site-packages/torch/_prims/utils.pygetnvFuserDtypesr.c seZdZUdZded<ded<edddddddd d d d d ddZeddddddfdd Zed ddZ ddZ ddZ Z S)! TensorMetaz Model tensor metadata. Not a stock meta tensor because device is modeled as the original device (not meta device), also we have different behavior for some high level Python bindings z Optional[Any]nodestrtnameN)shapestridesrdevicez5Optional[Union[TensorMeta, NumberType, torch.Tensor]]zOptional[ShapeType]zOptional[StrideType]Optional[torch.dtype]z"Optional[Union[torch.device, str]]) tensorlikerrrrc CsZt|tr^|s |dus$t|ts$J|s:|dus>t|ts>Jd}d}tt|}td} np|durt|ttjfszJt |j }t | }|j }|j} n0|dusJ|dusJ|dusJ|dusJ|dur|nt |}|dur|nt |}|dur|n|}|dur| n|}t|t r0t|}tjj|||d||dd} d| _d| _| S)NrcpurF)rstorage_offsetrr requires_grad) isinstanceNumberr type_to_dtypetypetorchrrTensortuplerstriderr_make_wrapper_subclassrr) clsrrrrrZinferred_shapeZinferred_stridesZinferred_dtypeZinferred_devicerrrr__new__8sF            zTensorMeta.__new__rrrz Sequence[Any]zOptional[Dict])functypesargskwargscs||dur i}|tjjjtjjtjjtjjjtjjjtjjjhvrRt ||||St |dslt d|d|j |i|S)Nmetaz Callable z has no meta function!)r'r(ndim__get__numelr*rrrsuper__torch_function__hasattr ValueErrorr3r,r/r0r1r2 __class__rrr8qs zTensorMeta.__torch_function__cCs tddS)Nzthis should be unreachable) RuntimeErrorr;rrr__torch_dispatch__szTensorMeta.__torch_dispatch__cCs|jSNr)selfrrr__repr__szTensorMeta.__repr__cCs|jSr@rA)rB format_specrrr __format__szTensorMeta.__format__)N)rN)rN) __name__ __module__ __qualname____doc____annotations__ staticmethodr. classmethodr8r?rCrE __classcell__rrr<rr,s* 8 rTensorLikeTypeabcCs t|tsJt|tsJt|j|jD](\}}||kr*d|j|j}t|q*|j|jkrxd|j|j}t||j|jkrt|jdkst|jdkrt|jdkst|jdkrnd|j|j}t|t ||\}}|sd| | |}t |dS)z Checks that two tensor likes have the same shape, dtype and device. In the future this will validate additional metadata, like strides. z!Shapes {0} and {1} are not equal!z!Dtypes {0} and {1} are not equal!zcuda:0cudaz"Devices {0} and {1} are not equal!z=Stride mismatch! Strides are {0} and {1} (mismatched at {2})!N) r# TensorLikeziprformatAssertionErrorrrrcheck_significant_stridesr*r>)rPrQxymsgZ same_stridesidxrrrcompare_tensor_metas0     r\zTuple[bool, Optional[int]])rPrQreturncCsj|jjdks|jjdkrf|dkrft|jD]6}||||kr.|j|dkr.d|fSq.dS)NrRrF)TN)rr&r6ranger4r*r)rPrQr[rrrrWs $&rWbool)rPr]cCsZ|dkrdSd}ttt|j|D](\}}|dkr>q,||krLdS||}q,dS)z Tests whether a tensor is contiguous or not. Tensors are contiguous when they have no elements, or when they have "nested" strides. rTr^F)r6reversedr)rTrr*)rPZexpected_striderXrYrrr is_contiguouss   rbzTuple[int, ...]r]c stdkrd}t|tdditddD}tddDtdkrZdSdj}|dkrpdS|d kr|d Sdjfd d }tt|}tt|t |dd }dg|}t |D]\}}|||<qt |} dg|} t |D]\}}| || |<qt| S)zA Computes the output strides for elementwise operations. rz:Can't compute elementwise output strides for zero tensors!allow_cpu_scalar_tensorsTcss|]}t|tVqdSr@)r#rS.0rPrrr z5compute_elementwise_output_strides..css$|]}t|trt|s|VqdSr@)r#rSis_cpu_scalar_tensorrerrrrgsrr^r^csD]v}||}||}|dks|dkr2q||kr@dS||krNdS||krddS||krdSqdS)Nrr^)r*)Zidx_aZidx_btensorZstride_aZstride_brtensorsrr_cmps  z0compute_elementwise_output_strides.._cmp)keyreverserk) lenr:check_same_shapeallr)r4rr_sortedr enumeratemake_contiguous_strides_for) rnrZZ all_tensorsr4ropermZpermuted_shaper[rX new_stridesZpermuted_stridesrrmr"compute_elementwise_output_stridess6       rzintlengthcCst|tsJ|dksJdS)zK Validates that an object represents a valid dimension length. rNr#r{r|rrrvalidate_dim_lengthJsr ShapeTypercCs$t|tsJ|D] }t|qdS)z= Validates that a sequence represents a valid shape. N)r#rr)rlrrrvalidate_shapeTsr StrideType)rcCs(t|tsJ|D]}|dksJqdS)z6 Verifies the object specifies valid strides. rN)r#r)rr*rrrvalidate_strides^srrankr[cCs<t|tsJt|tsJ|dkr,||ks8|dks8JdS)zr Validates that idx is a valid index for the given shape. Assumes the index is already canonicalized. rNr~rrrr validate_idxhsrDimsSequenceTyperindicescCs|D]}t||qdSr@)r)rrr[rrrvalidate_dimension_indicestsrrZex_idxcCs4t|tsJt|tsJ|dkr,||ks0JdS)zS Validates that ex_idx is a valid exclusive index for the given shape. rNr~rrrrvalidate_exclusive_idxysr)rr[r]cCsb|dkr |nd}|dkr$||kr$|S|dkr6||}n|}|dksJ||kr^d||}t||S)Nrr^z8Received out of bounds index {0} for tensor of rank {1}!)rUr:)rr[_rank_idxrZrrrcanonicalize_dims rDimsType)rrr]cs*t|trt|Stfdd|DS)Nc3s|]}t|VqdSr@)rrfrXrrrrgrhz$canonicalize_dims..)r#r{rr)rrrrcanonicalize_dimss  r)rrxr]cCs0t|tsdStt|ttd|ks,dSdS)z> Validates that perm is a permutation of length rank. FrT)r#rr)rur_)rrxrrris_valid_permutations  rrcCst|t|kS)z Compares two shapes a and b, returning True if they are the same (their ranks and corresponding lengths match) and False otherwise. )r)rOrrr is_same_shapesrrcCs t|to|jdko|jjdkS)Nrr)r#rSr4rr&)rPrrrrisricGst|dkrdSd}|D]}t|tr*qqt|tr|rBt|rBq|durP|j}||jkrdt|jdt|d}t|qdtt|d}t|qdS)a% Checks that all Tensors in args have the same device. Raises a RuntimeError when: - args contains an object whose type is not Tensor or Number - two Tensor objects in args have different devices, unless one is a CPU scalar tensor and allow_cpu_scalar_tensors is True r^NzTensor on device z is not on the expected device !z/Unexpected type when checking for same device, ) rrr#r$rSrirrr>r&)rdr1rargrZrrrcheck_same_devices4      rcGsd}|D]x}t|trqqt|trd|r2t|r2q|dur@|j}t||jsd|j|}t|qdtt |d}t|qdS)z Checks that all Tensors in args have the same shape. Raises a RuntimeError when: - args contains an object whose type is not Tensor or Number - two Tensor objects in args have different devices Nz(Shape {0} is not the expected shape {1}!z.Unexpected type when checking for same shape, r) r#r$rSrirrrUr>rr&)rdr1rrrZrrrrss"     rs)rr]cCst|tjsJ|tjuSr@)r#r'rr`rrrris_boolean_dtypesrcCst|tjsJ|tvSr@)r#r'r_integer_dtypesrrrris_integer_dtypesrcCst|tjsJ|tvSr@)r#r'r _float_dtypesrrrris_float_dtypesrcCst|tjsJ|tvSr@)r#r'r_complex_dtypesrrrris_complex_dtypesrcCst|Sr@)_complex_to_real_dtype_maprrrrcorresponding_real_dtype1srcCst|Sr@)_real_to_complex_dtype_maprrrrcorresponding_complex_dtype5srr&cCsNt|tjsJ|tjurtS|tvr*tS|tvr6tS|tvrBt St ddS)z[ Computes the corresponding Python type (AKA "type kind") for the given dtype. zInvalid dtype!N) r#r'rr`rr{rfloatrcomplexr:rrrr dtype_to_type9s r)typr]cCst|S)z= Computes the corresponding dtype for a Number type. )_type_to_dtype_map)rrrrr%Tsr%cCsZ|tvs J|tvsJ||ur$|StD]$}||ur<|S||ur(|Sq(tddS)zw Returns the higher of the two given Number types. The types are ordered bool -> int -> float -> complex. zUnknown Python scalar type!N)_ordered_typesr:)rPrQrrrrget_higher_type^s   r8Optional[Union[torch.dtype, TensorLikeType, NumberType]]rc Cs8|dust|tjttfsJ|dusrXrrr_extract_dtypes    z(get_higher_dtype.._extract_dtyper^rUnexpected termination!)r#r'rrSr$r`uint8int8int16int32int64float16bfloat16float32float64 complex32 complex64 complex128rvr>)rPrQrZordered_datatypesr[dtypesrrrget_higher_dtypews<     rcCsVttttf}||vsJ||vs$J|D] }||kr:dS||kr(dSq(tddS)z Compares two types, a and b, returning True if a is weakly "less" than b. The comparison is determined by the following type ordering: bool, int, float, complex. TFrN)r`r{rrr>)rPrQZ ordered_typesrrrris_weakly_lesser_types  r)cast_to cast_fromr]cCsBttttfD] }||rdS||r dSq td||dS)NTFz!Received unknown dtypes {0}, {1}!)rrrrr:rU)rrfnrrrcan_safe_cast_tos rcGsd}d}|D]}t|trq q t|tr|dur6|j}|durHt|j}||jurxdt|jdt|d}t|t|j}||urdt|dt|d}t|q dtt|d}t|q dS)a Checks that all Tensors in args have the same device and that all Numbers have the same corresponding Python type. Raises a RuntimeError when: - args contains an object whose type is not Tensor or Number - two Tensors objects in args have different dtypes - two Number objects in args have different types - there are Tensors and Numbers in args, and one of those Tensors corresponding Python types is different from the type of one of those Numbers NzTensor with dtype z is not the expected dtype of rz&Tensor with corresponding Python type z is not the expected type of z.Unexpected type when checking for same dtype, )r#r$rSrrrr>r&)r1 full_dtype scalar_typerrZarg_typerrrcheck_same_dtypesL       rcCs t||Sr@)_computation_dtype_mapgetrrrr_get_computation_dtype"src@s$eZdZdZdZdZdZdZdZdS)ELEMENTWISE_TYPE_PROMOTION_KIND)rrj))))N) rFrGrHDEFAULT NO_OPMATH INT_TO_FLOAT ALWAYS_BOOLCOMPLEX_TO_FLOAT BOOL_TO_LONGrrrrr&s rzTuple[torch.dtype, torch.dtype])type_promotion_kindr]c Gstdd|D}t}|D]V}t|ttfsFdtt|}t|t|tr`t |t|}qt |t |j }qd}ddddd d }|t ur||t }|durtn|}n`|tur||d d d d}|durtt}n0|tur ||t}|durtjn|}ntj}|tjur*t||fS|tjur>||fS|tjurrt|s^t|rft}t||fS|tjurt|}t|rt|}||fS|tjurt|rtjtjfSt||fS|tjurt|tjfStdt|dS)a Computes the computation and result dtypes for elementwise type promotion on the given arguments and with the given elementwise type promotion kind. Note that not all inputs to an elementwise operation necessarily participate in type promotion. For example, the "alpha" parameter of torch.add does not participate in type promotion, although it may be cast to the Python type corresponding to the computation dtype that the type promotion algorithm determines. Default elementwise type promotion, which all other type promotion kinds tweak (see below), first decides which of four ordered types to use: bool -> integer -> floating point -> complex The selected type is the "lowest" type in the above list such that all number arguments have a weakly "lower" type and all tensor arguments have a weakly lower corresponding type for their dtype. Once the type is determined, the particular result dtype is found. The dtypes are partially ordered as follows: bool -> uint8, int8 -> int16 -> int32 -> int64 -> float16, bfloat16 -> float32 -> float64 -> complex32 -> complex64 -> complex128 The result dtype is selected by: - if no tensor's dtype has the same corresponding type as the one selected, then the result dtype is the (default) dtype corresponding to the selected type (for example, 1.5 + an integer tensor has a result dtype of the default floating point dtype) - if the result type is complex then the dtype is: - the default complex dtype if there are no floating point or complex tensors - if there are floating point or complex tensors with one or more dimensions, then the complex dtype corresponding to the highest corresponding complex dtype among those tensors (for example, double + cfloat -> cdouble) - if there are only floating point or complex tensors with zero dimensions, then the complex dtype corresponding to the highest corresponding complex dtype among those tensors - if the first two cases do not apply, the result dtype is the highest dtype among all tensors with one or more dimensions of the output type, and if there are no such tensors then it's the highest dtype among all tensors with zero dimensions of the output type (for example, long + half -> half, even if the half tensor has zero dimensions) The "corresponding complex dtypes" are: float16 -> complex32 bfloat16 -> complex64 float32 -> complex64 float64 -> complex128 complex32 -> complex32 complex64 -> complex64 complex128 -> complex128 The DEFAULT type promotion kind computes per above, and then uses the result dtype to pick a computation dtype by mapping low precision floating point and complex dtypes as follows: float16 -> float32 bfloat16 -> float32 complex32 -> complex64 This is referred to as "op math", and the NO_OPMATH type promotion kind disables this mapping, making the computation dtype the same as the result dtype when it's selected. NO_OPMATH is appropriate for kernels which perform no mathematical operations on their tensors (see below for examples). The INT_TO_FLOAT type promotion kind maps boolean and integer maps result dtypes to the default floating point dtype, and computation dtypes to the appropriate op math dtype. The COMPLEX_TO_FLOAT type promotion kind maps complex result dtypes to the corresponding float dtype, following this mapping: complex32 -> float16 complex64 -> float32 complex128 -> float64 Note that COMPLEX_TO_FLOAT derives the computation dtype as the DEFAULT setting does. The BOOL_TO_LONG type promotion kind maps boolean computation and result dtypes to long. The ALWAYS_BOOL type promotion kind always sets the result dtype to bool. Example operators for each type promotion option: DEFAULT : add NO_OPMATH : where, nextafter, cat INT_TO_FLOAT : sin COMPLEX_TO_FLOAT : abs BOOL_TO_LONG : pow ALWAYS_BOOL : eq css|]}|dur|VqdSr@rrrrrrgrhz%elementwise_dtypes..z>Unexpected type {0} when computing elementwise type promotion!NF)float_as_complexrrccSspd}d}|D]R}t|tr ||jr |j}|r>t|r>t|}|jdkrTt||}q t||}q |durl|S|S)Nr)r#rSrrrr4r)r1filterrZzero_dim_tensor_dtypeZone_plus_dim_tensor_dtyperX_dtyperrr_find_highest_dtype_filtereds"  z8elementwise_dtypes.._find_highest_dtype_filteredcSst|pt|Sr@)rrrrrrrhz$elementwise_dtypes..TzUnknown type promotion kind {0}) r)r`r#r$rSrUrr&r:rrrrrr'get_default_dtyperrr{rlongrrrrrrrrrrr) r_argsr1Z highest_typerXrZ result_dtyperZcomputation_dtyperrrelementwise_dtypes0sjZ                  rzUnion[str, torch.device]z torch.device)dr]cCs,t|ttjfsJt|tr(t|S|S)zk Wraps strings into torch.device objects. Given torch.device objects are returned unmodified. )r#rr'r)rrrr wrap_devices  r)rr]cCs\t||sdSd}g}t|D]*}|dkr@||||}q ||q tt|}|S)Nrr^r)rraappendr))rZ multiplierrrresultrrrrws     rw)r dimensionsr]cCsN|D]}tt||qg}tt|D]}||vr6q(|||q(t|Sr@)rrrr_rr))rrr[ new_shaperrrcompute_reduction_output_shapesrdimscCs t|tt|krtddS)Nz#duplicate value in the list of dims)rrsetr>rrrrvalidate_no_repeating_dimssrzOptional[Sequence])rrr]cs:|durtttStfdd|D}t||S)Nc3s|]}tt|VqdSr@)rrr)rfr[rrrrgrhz!reduction_dims..)r)r_rrr)rrrrrreduction_dimss rztorch._TypedStorage)rPrrr c Csttj|dkrdS||}d}t||D]\}}||d|}q.||kr~||}d||t|t||} t| dS)z] Determines if the given shape, strides, and offset are valid for the given storage. rNr^z~Can't view a storage of size {0} with an offset of {1}, shape of {2}, and strides of {3}, which requires a storage of size {4})r operatormulsizerTrUrr:) rPrrr r}Z max_offsetrXrYZrequired_lengthrZrrrcheck_in_bounds_for_storage#s r)q __future__rtypingrrrrrrr r enumr functoolsr r rr'r9_CZtorch._C._nvfuserrcdouble ComplexDoublecfloat ComplexFloatdoubleDoublerFloathalfHalfrBFloat16rIntr{Int32r`BoolrrSizerrrrr NumberTyper$r(rrNrSZTensorSequenceTyper\rWrbrzrrrrrrrrrrrirrsrrrrrrrrrrrrrrrrrrrrrrrrr%rrrrrrrrrrrrwrrrrrrrrs (     r 'P      +  @ B 7