^jXddlmZddlmZddlZGddeeZeeefZdefdZdefdZ d ejdefd Z d ejdefd Z y) )Enum)UnionNceZdZdZdZdZdZy)FormatNCHWNHWCNCLNLCN)__name__ __module__ __qualname__rrr r ]/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/timm/layers/format.pyrrs D D C Crrfmtct|}|tjurd}|S|tjurd}|S|tjurd}|Sd}|S)zReturn spatial dimension indices for a given tensor format. Args: fmt: Tensor format (NCHW, NHWC, NCL, or NLC). Returns: Tuple of spatial dimension indices. )))rr)r)rr r rrdims rget_spatial_dimrse +C fjj J   J   J Jrcxt|}|tjurd}|S|tjurd}|Sd}|S)zReturn channel dimension index for a given tensor format. Args: fmt: Tensor format (NCHW, NHWC, NCL, or NLC). Returns: Channel dimension index. rrr)rrr rs rget_channel_dimr&sK +C fkk J   J Jrxc |tjk(r|jdddd}|S|tjk(r#|j dj dd}|S|tj k(r|j d}|S)zConvert tensor from NCHW format to specified format. Args: x: Input tensor in NCHW format. fmt: Target format. Returns: Tensor in target format. rrrr)rrpermuter flatten transposer rrs rnchw_tor!9sz fkk IIaAq ! H   IIaL " "1a ( H   IIaL Hrc|tjk(r|jdddd}|S|tjk(r|j dd}|S|tj k(r"|j ddj dd}|S)zConvert tensor from NHWC format to specified format. Args: x: Input tensor in NHWC format. fmt: Target format. Returns: Tensor in target format. rrrr)rrrr rr rr s rnhwc_tor#Ls~ fkk IIaAq ! H   IIaO H   IIaO % %a + Hr) enumrtypingrtorchstrrFormatTrrTensorr!r#rrrr*sr S$ V *& u|| & & u|| & r