a  þdïã@sdddlZddlmZGdd„dejƒZejjGdd„dƒƒZGdd„dejƒZGdd „d ejƒZ dS) éNcs8eZdZUejjeed<d‡fdd„ Zdd„Z ‡Z S)Ú SpaceToDepthÚbsécs tƒ ¡|dksJ‚||_dS)Nr©ÚsuperÚ__init__r©ÚselfÚ block_size©Ú __class__©úc/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/timm/layers/space_to_depth.pyrs  zSpaceToDepth.__init__cCsx| ¡\}}}}| ||||j|j||j|j¡}| dddddd¡ ¡}| |||j|j||j||j¡}|S)Nréééér©ÚsizeÚviewrÚpermuteÚ contiguous©r ÚxÚNÚCÚHÚWr r rÚforward s $(zSpaceToDepth.forward)r) Ú__name__Ú __module__Ú __qualname__ÚtorchÚjitÚFinalÚintÚ__annotations__rrÚ __classcell__r r r rrs rc@seZdZejdœdd„ZdS)ÚSpaceToDepthJit)rcCsd| ¡\}}}}| |||dd|dd¡}| dddddd¡ ¡}| ||d|d|d¡}|S)Nrrrrrré)rrrrrr r rÚ__call__s zSpaceToDepthJit.__call__N)rr r!r"ZTensorr*r r r rr(sr(cs&eZdZd‡fdd„ Zdd„Z‡ZS)ÚSpaceToDepthModuleFcs$tƒ ¡|stƒ|_ntƒ|_dS©N)rrr(Úopr)r Zno_jitr r rr!s  zSpaceToDepthModule.__init__cCs | |¡Sr,)r-)r rr r rr(szSpaceToDepthModule.forward)F©rr r!rrr'r r r rr+ sr+cs$eZdZ‡fdd„Zdd„Z‡ZS)Ú DepthToSpacecstƒ ¡||_dSr,rrr r rr.s zDepthToSpace.__init__cCst| ¡\}}}}| ||j|j||jd||¡}| dddddd¡ ¡}| |||jd||j||j¡}|S)Nrrrrrrrrr r rr2s "&zDepthToSpace.forwardr.r r r rr/,s r/) r"Ztorch.nnÚnnÚModulerr#Úscriptr(r+r/r r r rÚs