o wÖiZã@shUddlmZddlZddlZddlmZmZmZmZm Z dZ e e d<Gdd„de ƒZGdd „d ƒZdS) é)Ú EinopsErrorN)ÚListÚOptionalÚSetÚTupleÚUnionuâ€¦Ú _ellipsisc@s&eZdZdZdefdd„Zdd„ZdS)Ú AnonymousAxiszHImportant thing: all instances of this class are not equal to each otherÚvaluecCs:t|ƒ|_|jdkr|jdkrtdƒ‚td |j¡ƒ‚dS)NézENo need to create anonymous axis of length 1. Report this as an issuez2Anonymous axis should have positive length, not {})Úintr rÚformat)Úselfr ©rúK/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/einops/parsing.pyÚ__init__ s   üzAnonymousAxis.__init__cCsd t|jƒ¡S)Nz{}-axis)r Ústrr )rrrrÚ__repr__szAnonymousAxis.__repr__N)Ú__name__Ú __module__Ú __qualname__Ú__doc__rrrrrrrr s r c @s„eZdZdZdddœdededefdd„Zd efd d „Zd efd d „Z e ddeded e eeffdd„ƒZ e ded efdd„ƒZ dS)ÚParsedExpressionz¡ non-mutable structure that contains information about one side of expression (e.g. 'b c (h w)') and keeps some information important for downstream F)Úallow_underscoreÚallow_duplicatesÚ expressionrrcsbdˆ_dˆ_tƒˆ_dˆ_gˆ_d|vr9d|vrtdƒ‚t |d¡dks,t |d¡dkr0tdƒ‚|  dt ¡}dˆ_d‰‡‡‡‡fd d „}d}|D]Q}|d vr||durV||ƒd}|d krgˆdurdtd ƒ‚g‰qH|dkr{ˆdurstdƒ‚ˆj  ˆ¡d‰qHt  |¡s‡|dt fvr“|durŽ|}qH||7}qHtd  |¡ƒ‚ˆdur¥td  |¡ƒ‚|dur¯||ƒdSdS)NFÚ.z...z6Expression may contain dots only inside ellipsis (...)r ézUExpression may contain dots only inside ellipsis (...); only one ellipsis for tensor Tcs|ˆjvrˆr |dksˆstd |¡ƒ‚|tkr7ˆj t¡ˆdur-ˆj t¡dˆ_dSˆ t¡dˆ_dSt  |¡}|rSt |ƒdkrSˆdurPˆj g¡dS dSˆj |ˆd\}}|sh|shtd ||¡ƒ‚|rnt |ƒ}ˆj |¡|rydˆ_ ˆdur†ˆj |g¡dSˆ |¡dS)NÚ_z5Indexing expression contains duplicate dimension "{}"FTr )rzInvalid axis identifier: {} {})Ú identifiersrr rÚaddÚ compositionÚappendÚhas_ellipsis_parenthesizedrÚ isdecimalr Úcheck_axis_name_return_reasonr Úhas_non_unitary_anonymous_axes)ÚxÚ is_numberÚ is_axis_nameÚreason©rrÚ bracket_grouprrrÚ add_axis_name2s8        ÿ z0ParsedExpression.__init__..add_axis_namez() ú(zDAxis composition is one-level (brackets inside brackets not allowed)ú)zBrackets are not balancedrzUnknown character '{}'z*Imbalanced parentheses in expression: "{}")Ú has_ellipsisr#Úsetrr&r!rrÚcountÚreplacerr"Úisalnumr )rrrrr-Úcurrent_identifierÚcharrr+rrsT ÿ " €  ÿzParsedExpression.__init__ÚreturncCs:g}|jD]}t|tƒsJdƒ‚|D]}| |¡qq|S)Nzdoes not work with ellipsis)r!Ú isinstanceÚlistr")rÚresultÚ composed_axisÚaxisrrrÚflat_axes_orderps  ÿz ParsedExpression.flat_axes_ordercCs,|jD]}t|tƒrt|ƒdkrdSqdS)Nr TF)r!r8r9Úlen)rÚaxesrrrÚhas_composed_axesxs €z"ParsedExpression.has_composed_axesÚnamecCsnt |¡sdS|ddks|ddkr|dkr|rdSdSt |¡r+t d |¡t¡|dvr5t d t¡dS) N)Fznot a valid python identifierrréÿÿÿÿ)TÚ)Fz8axis name should should not start or end with underscorez9It is discouraged to use axes names that are keywords: {})r<zQIt is discouraged to use 'axis' as an axis name and will raise an error in future) rÚ isidentifierÚkeywordÚ iskeywordÚwarningsÚwarnr ÚRuntimeWarningÚ FutureWarning)rArrrrr%s   þz.ParsedExpression.check_axis_name_return_reasoncCst |¡\}}|S)z‹ Valid axes names are python identifiers except keywords, and additionally should not start or end with underscore )rr%)rAÚis_validÚ_reasonrrrÚcheck_axis_name‘sz ParsedExpression.check_axis_nameN)F)rrrrrÚboolrrr=r@Ú staticmethodrr%rMrrrrrsR"r)ÚeinopsrrErGÚtypingrrrrrrrÚ__annotations__Úobjectr rrrrrÚs