





'''simple docstring'''







import argparse
import json
import math
import os
import time
import traceback
import zipfile
from collections import Counter

import requests





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	Tuple,	SCREAMING_SNAKE_CASE__					:	List[Any]=None      )       ->		Optional[int]:
						UpperCAmelCase_	:					Tuple     =     None
						if token is not None:
												UpperCAmelCase_	:					Any     =     {'''Accept''': '''application/vnd.github+json''', '''Authorization''': F"""Bearer {token}"""}

						UpperCAmelCase_	:					List[Any]     =     F"""https://api.github.com/repos/huggingface/transformers/actions/runs/{workflow_run_id}/jobs?per_page=100"""
						UpperCAmelCase_	:					int     =     requests.get(A_,	headers=A_      ).json()
						UpperCAmelCase_	:					Optional[int]     =     {}

						try:
												job_links.update({job['''name''']: job['''html_url'''] for job in result['''jobs''']}      )
												UpperCAmelCase_	:					str     =     math.ceil((result['''total_count'''] - 100) / 100      )

												for i in range(A_      ):
																		UpperCAmelCase_	:					Tuple     =     requests.get(url + F"""&page={i + 2}""",	headers=A_      ).json()
																		job_links.update({job['''name''']: job['''html_url'''] for job in result['''jobs''']}      )

												return job_links
						except Exception:
												print(F"""Unknown error, could not fetch links:\n{traceback.format_exc()}"""      )

						return {}





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	str,	SCREAMING_SNAKE_CASE__					:	Union[str, Any]=None      )       ->		int:
						UpperCAmelCase_	:					Optional[int]     =     None
						if token is not None:
												UpperCAmelCase_	:					List[str]     =     {'''Accept''': '''application/vnd.github+json''', '''Authorization''': F"""Bearer {token}"""}

						UpperCAmelCase_	:					Optional[int]     =     F"""https://api.github.com/repos/huggingface/transformers/actions/runs/{worflow_run_id}/artifacts?per_page=100"""
						UpperCAmelCase_	:					Optional[int]     =     requests.get(A_,	headers=A_      ).json()
						UpperCAmelCase_	:					Optional[Any]     =     {}

						try:
												artifacts.update({artifact['''name''']: artifact['''archive_download_url'''] for artifact in result['''artifacts''']}      )
												UpperCAmelCase_	:					Optional[Any]     =     math.ceil((result['''total_count'''] - 100) / 100      )

												for i in range(A_      ):
																		UpperCAmelCase_	:					int     =     requests.get(url + F"""&page={i + 2}""",	headers=A_      ).json()
																		artifacts.update({artifact['''name''']: artifact['''archive_download_url'''] for artifact in result['''artifacts''']}      )

												return artifacts
						except Exception:
												print(F"""Unknown error, could not fetch links:\n{traceback.format_exc()}"""      )

						return {}





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	Optional[Any],	SCREAMING_SNAKE_CASE__					:	List[str],	SCREAMING_SNAKE_CASE__					:	int,	SCREAMING_SNAKE_CASE__					:	str      )       ->		Optional[Any]:
						UpperCAmelCase_	:					Dict     =     None
						if token is not None:
												UpperCAmelCase_	:					Any     =     {'''Accept''': '''application/vnd.github+json''', '''Authorization''': F"""Bearer {token}"""}

						UpperCAmelCase_	:					List[Any]     =     requests.get(A_,	headers=A_,	allow_redirects=A_      )
						UpperCAmelCase_	:					str     =     result.headers['''Location''']
						UpperCAmelCase_	:					Union[str, Any]     =     requests.get(A_,	allow_redirects=A_      )
						UpperCAmelCase_	:					Dict     =     os.path.join(A_,	F"""{artifact_name}.zip"""      )
						with open(A_,	'''wb'''      ) as fp:
												fp.write(response.content      )





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	List[str],	SCREAMING_SNAKE_CASE__					:	Optional[int]=None      )       ->		Union[str, Any]:
						UpperCAmelCase_	:					int     =     []
						UpperCAmelCase_	:					Union[str, Any]     =     []
						UpperCAmelCase_	:					Optional[int]     =     None

						with zipfile.ZipFile(A_      ) as z:
												for filename in z.namelist():
																		if not os.path.isdir(A_      ):
																								# read the file
																								if filename in ["failures_line.txt", "summary_short.txt", "job_name.txt"]:
																														with z.open(A_      ) as f:
																																				for line in f:
																																										UpperCAmelCase_	:					Dict     =     line.decode('''UTF-8'''      ).strip()
																																										if filename == "failures_line.txt":
																																																try:
																																																						# `error_line` is the place where `error` occurs
																																																						UpperCAmelCase_	:					str     =     line[: line.index(''': '''      )]
																																																						UpperCAmelCase_	:					Tuple     =     line[line.index(''': '''      ) + len(''': '''      ) :]
																																																						errors.append([error_line, error]      )
																																																except Exception:
																																																						# skip un-related lines
																																																						pass
																																										elif filename == "summary_short.txt" and line.startswith('''FAILED '''      ):
																																																# `test` is the test method that failed
																																																UpperCAmelCase_	:					Optional[Any]     =     line[len('''FAILED '''      ) :]
																																																failed_tests.append(A_      )
																																										elif filename == "job_name.txt":
																																																UpperCAmelCase_	:					Union[str, Any]     =     line

						if len(A_      ) != len(A_      ):
												raise ValueError(
												    F"""`errors` and `failed_tests` should have the same number of elements. Got {len(A_      )} for `errors` """
												    F"""and {len(A_      )} for `failed_tests` instead. The test reports in {artifact_zip_path} have some"""
												    ''' problem.'''      )

						UpperCAmelCase_	:					Optional[int]     =     None
						if job_name and job_links:
												UpperCAmelCase_	:					int     =     job_links.get(A_,	A_      )

						# A list with elements of the form (line of error, error, failed test)
						UpperCAmelCase_	:					Optional[int]     =     [x + [y] + [job_link] for x, y in zip(A_,	A_      )]

						return result





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	int,	SCREAMING_SNAKE_CASE__					:	Dict=None      )       ->		str:
						UpperCAmelCase_	:					Union[str, Any]     =     []

						UpperCAmelCase_	:					List[str]     =     [os.path.join(A_,	A_      ) for p in os.listdir(A_      ) if p.endswith('''.zip'''      )]
						for p in paths:
												errors.extend(get_errors_from_single_artifact(A_,	job_links=A_      )      )

						return errors





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	Union[str, Any],	SCREAMING_SNAKE_CASE__					:	List[str]=None      )       ->		int:
						UpperCAmelCase_	:					Tuple     =     Counter()
						counter.update([x[1] for x in logs]      )
						UpperCAmelCase_	:					List[str]     =     counter.most_common()
						UpperCAmelCase_	:					str     =     {}
						for error, count in counts:
												if error_filter is None or error not in error_filter:
																		UpperCAmelCase_	:					List[str]     =     {'''count''': count, '''failed_tests''': [(x[2], x[0]) for x in logs if x[1] == error]}

						UpperCAmelCase_	:					Dict     =     dict(sorted(r.items(),	key=lambda SCREAMING_SNAKE_CASE__      : item[1]["count"],	reverse=A_      )      )
						return r





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	str      )       ->		Union[str, Any]:
						UpperCAmelCase_	:					Dict     =     test.split('''::'''      )[0]
						if test.startswith('''tests/models/'''      ):
												UpperCAmelCase_	:					Union[str, Any]     =     test.split('''/'''      )[2]
						else:
												UpperCAmelCase_	:					Any     =     None

						return test





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	List[Any],	SCREAMING_SNAKE_CASE__					:	Optional[Any]=None      )       ->		Any:
						UpperCAmelCase_	:					Tuple     =     [(x[0], x[1], get_model(x[2]      )) for x in logs]
						UpperCAmelCase_	:					List[str]     =     [x for x in logs if x[2] is not None]
						UpperCAmelCase_	:					int     =     {x[2] for x in logs}

						UpperCAmelCase_	:					List[str]     =     {}
						for test in tests:
												UpperCAmelCase_	:					Tuple     =     Counter()
												# count by errors in `test`
												counter.update([x[1] for x in logs if x[2] == test]      )
												UpperCAmelCase_	:					List[str]     =     counter.most_common()
												UpperCAmelCase_	:					List[Any]     =     {error: count for error, count in counts if (error_filter is None or error not in error_filter)}
												UpperCAmelCase_	:					List[Any]     =     sum(error_counts.values()      )
												if n_errors > 0:
																		UpperCAmelCase_	:					List[str]     =     {'''count''': n_errors, '''errors''': error_counts}

						UpperCAmelCase_	:					str     =     dict(sorted(r.items(),	key=lambda SCREAMING_SNAKE_CASE__      : item[1]["count"],	reverse=A_      )      )
						return r





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	Dict      )       ->		str:
						UpperCAmelCase_	:					Optional[Any]     =     '''| no. | error | status |'''
						UpperCAmelCase_	:					Tuple     =     '''|-:|:-|:-|'''
						UpperCAmelCase_	:					List[Any]     =     [header, sep]
						for error in reduced_by_error:
												UpperCAmelCase_	:					Tuple     =     reduced_by_error[error]['''count''']
												UpperCAmelCase_	:					Dict     =     F"""| {count} | {error[:100]} |  |"""
												lines.append(A_      )

						return "\n".join(A_      )





def 				lowerCamelCase_					(     SCREAMING_SNAKE_CASE__					:	Union[str, Any]      )       ->		int:
						UpperCAmelCase_	:					Any     =     '''| model | no. of errors | major error | count |'''
						UpperCAmelCase_	:					Dict     =     '''|-:|-:|-:|-:|'''
						UpperCAmelCase_	:					Union[str, Any]     =     [header, sep]
						for model in reduced_by_model:
												UpperCAmelCase_	:					int     =     reduced_by_model[model]['''count''']
												UpperCAmelCase_  ,				UpperCAmelCase_	:					Tuple     =     list(reduced_by_model[model]['''errors'''].items()      )[0]
												UpperCAmelCase_	:					str     =     F"""| {model} | {count} | {error[:60]} | {_count} |"""
												lines.append(A_      )

						return "\n".join(A_      )


if __name__ == "__main__":
					snake_case_   :				int		=     argparse.ArgumentParser()
					# Required parameters
					parser.add_argument("--workflow_run_id", type=str, required=True, help="A GitHub Actions workflow run id.")
					parser.add_argument(
					    "--output_dir",
					    type=str,
					    required=True,
					    help="Where to store the downloaded artifacts and other result files.",
					)
					parser.add_argument("--token", default=None, type=str, help="A token that has actions:read permission.")
					snake_case_   :				List[Any]		=     parser.parse_args()

					os.makedirs(args.output_dir, exist_ok=True)

					snake_case_   :				List[Any]		=     get_job_links(args.workflow_run_id, token=args.token)
					snake_case_   :				Tuple		=     {}
					# To deal with `workflow_call` event, where a job name is the combination of the job names in the caller and callee.
					# For example, `PyTorch 1.11 / Model tests (models/albert, single-gpu)`.
					if _job_links:
										for k, v in _job_links.items():
															# This is how GitHub actions combine job names.
															if " / " in k:
																				snake_case_   :				List[Any]		=     k.find(" / ")
																				snake_case_   :				Optional[Any]		=     k[index + len(" / ") :]
															snake_case_   :				Optional[Any]		=     v
					with open(os.path.join(args.output_dir, "job_links.json"), "w", encoding="UTF-8") as fp:
										json.dump(job_links, fp, ensure_ascii=False, indent=4)

					snake_case_   :				Dict		=     get_artifacts_links(args.workflow_run_id, token=args.token)
					with open(os.path.join(args.output_dir, "artifacts.json"), "w", encoding="UTF-8") as fp:
										json.dump(artifacts, fp, ensure_ascii=False, indent=4)

					for idx, (name, url) in enumerate(artifacts.items()):
										download_artifact(name, url, args.output_dir, args.token)
										# Be gentle to GitHub
										time.sleep(1)

					snake_case_   :				str		=     get_all_errors(args.output_dir, job_links=job_links)

					# `e[1]` is the error
					snake_case_   :				List[str]		=     Counter()
					counter.update([e[1] for e in errors])

					# print the top 30 most common test errors
					snake_case_   :				List[Any]		=     counter.most_common(30)
					for item in most_common:
										print(item)

					with open(os.path.join(args.output_dir, "errors.json"), "w", encoding="UTF-8") as fp:
										json.dump(errors, fp, ensure_ascii=False, indent=4)

					snake_case_   :				int		=     reduce_by_error(errors)
					snake_case_   :				Union[str, Any]		=     reduce_by_model(errors)

					snake_case_   :				Dict		=     make_github_table(reduced_by_error)
					snake_case_   :				Any		=     make_github_table_per_model(reduced_by_model)

					with open(os.path.join(args.output_dir, "reduced_by_error.txt"), "w", encoding="UTF-8") as fp:
										fp.write(sa)
					with open(os.path.join(args.output_dir, "reduced_by_model.txt"), "w", encoding="UTF-8") as fp:
										fp.write(sa)







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block




@dataclass
class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  a__           =		42







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: Tuple			,							UpperCamelCase__			: List[Any]=3			,							UpperCamelCase__			: str=3			,							UpperCamelCase__			: List[Any]=("DownEncoderBlock2D",)			,							UpperCamelCase__			: Optional[Any]=(64,)			,							UpperCamelCase__			: Optional[Any]=2			,							UpperCamelCase__			: Union[str, Any]=32			,							UpperCamelCase__			: Optional[Any]="silu"			,							UpperCamelCase__			: List[str]=True			,							)			->	str:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       layers_per_block

    __magic_name__				     =       torch.nn.Convad(
        UpperCamelCase__			,							block_out_channels[0]			,							kernel_size=3			,							stride=1			,							padding=1			,							)

    __magic_name__				     =       None
    __magic_name__				     =       nn.ModuleList([]    )

    # down
    __magic_name__				     =       block_out_channels[0]
    for i, down_block_type in enumerate(UpperCamelCase__    ):
      __magic_name__				     =       output_channel
      __magic_name__				     =       block_out_channels[i]
      __magic_name__				     =       i == len(UpperCamelCase__    ) - 1

      __magic_name__				     =       get_down_block(
          UpperCamelCase__			,							num_layers=self.layers_per_block			,							in_channels=UpperCamelCase__			,							out_channels=UpperCamelCase__			,							add_downsample=not is_final_block			,							resnet_eps=1E-6			,							downsample_padding=0			,							resnet_act_fn=UpperCamelCase__			,							resnet_groups=UpperCamelCase__			,							attention_head_dim=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							)
      self.down_blocks.append(UpperCamelCase__    )

    # mid
    __magic_name__				     =       UNetMidBlockaD(
        in_channels=block_out_channels[-1]			,							resnet_eps=1E-6			,							resnet_act_fn=UpperCamelCase__			,							output_scale_factor=1			,							resnet_time_scale_shift="""default"""			,							attention_head_dim=block_out_channels[-1]			,							resnet_groups=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							)

    # out
    __magic_name__				     =       nn.GroupNorm(num_channels=block_out_channels[-1]			,							num_groups=UpperCamelCase__			,							eps=1E-6    )
    __magic_name__				     =       nn.SiLU()

    __magic_name__				     =       2 * out_channels if double_z else out_channels
    __magic_name__				     =       nn.Convad(block_out_channels[-1]			,							UpperCamelCase__			,							3			,							padding=1    )

    __magic_name__				     =       False





  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: Optional[Any]    )			->	int:







    """simple docstring"""






    __magic_name__				     =       x
    __magic_name__				     =       self.conv_in(UpperCamelCase__    )

    if self.training and self.gradient_checkpointing:

      def create_custom_forward(UpperCamelCase__			: int    ):
        def custom_forward(*UpperCamelCase__			: str    ):
          return module(*UpperCamelCase__    )

        return custom_forward

      # down
      if is_torch_version(""">="""			,							"""1.11.0"""    ):
        for down_block in self.down_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(
              create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block    )			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
      else:
        for down_block in self.down_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__    )
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block    )			,							UpperCamelCase__    )

    else:
      # down
      for down_block in self.down_blocks:
        __magic_name__				     =       down_block(UpperCamelCase__    )

      # middle
      __magic_name__				     =       self.mid_block(UpperCamelCase__    )

    # post-process
    __magic_name__				     =       self.conv_norm_out(UpperCamelCase__    )
    __magic_name__				     =       self.conv_act(UpperCamelCase__    )
    __magic_name__				     =       self.conv_out(UpperCamelCase__    )

    return sample







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: List[str]			,							UpperCamelCase__			: int=3			,							UpperCamelCase__			: Dict=3			,							UpperCamelCase__			: List[Any]=("UpDecoderBlock2D",)			,							UpperCamelCase__			: List[Any]=(64,)			,							UpperCamelCase__			: Dict=2			,							UpperCamelCase__			: int=32			,							UpperCamelCase__			: Optional[int]="silu"			,							UpperCamelCase__			: Tuple="group"			,							)			->	Dict:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       layers_per_block

    __magic_name__				     =       nn.Convad(
        UpperCamelCase__			,							block_out_channels[-1]			,							kernel_size=3			,							stride=1			,							padding=1			,							)

    __magic_name__				     =       None
    __magic_name__				     =       nn.ModuleList([]    )

    __magic_name__				     =       in_channels if norm_type == """spatial""" else None

    # mid
    __magic_name__				     =       UNetMidBlockaD(
        in_channels=block_out_channels[-1]			,							resnet_eps=1E-6			,							resnet_act_fn=UpperCamelCase__			,							output_scale_factor=1			,							resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type			,							attention_head_dim=block_out_channels[-1]			,							resnet_groups=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							)

    # up
    __magic_name__				     =       list(reversed(UpperCamelCase__    )    )
    __magic_name__				     =       reversed_block_out_channels[0]
    for i, up_block_type in enumerate(UpperCamelCase__    ):
      __magic_name__				     =       output_channel
      __magic_name__				     =       reversed_block_out_channels[i]

      __magic_name__				     =       i == len(UpperCamelCase__    ) - 1

      __magic_name__				     =       get_up_block(
          UpperCamelCase__			,							num_layers=self.layers_per_block + 1			,							in_channels=UpperCamelCase__			,							out_channels=UpperCamelCase__			,							prev_output_channel=UpperCamelCase__			,							add_upsample=not is_final_block			,							resnet_eps=1E-6			,							resnet_act_fn=UpperCamelCase__			,							resnet_groups=UpperCamelCase__			,							attention_head_dim=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							resnet_time_scale_shift=UpperCamelCase__			,							)
      self.up_blocks.append(UpperCamelCase__    )
      __magic_name__				     =       output_channel

    # out
    if norm_type == "spatial":
      __magic_name__				     =       SpatialNorm(block_out_channels[0]			,							UpperCamelCase__    )
    else:
      __magic_name__				     =       nn.GroupNorm(num_channels=block_out_channels[0]			,							num_groups=UpperCamelCase__			,							eps=1E-6    )
    __magic_name__				     =       nn.SiLU()
    __magic_name__				     =       nn.Convad(block_out_channels[0]			,							UpperCamelCase__			,							3			,							padding=1    )

    __magic_name__				     =       False





  def 							_lowercase							( self			: str			,							UpperCamelCase__			: Union[str, Any]			,							UpperCamelCase__			: Tuple=None    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       z
    __magic_name__				     =       self.conv_in(UpperCamelCase__    )

    __magic_name__				     =       next(iter(self.up_blocks.parameters()    )    ).dtype
    if self.training and self.gradient_checkpointing:

      def create_custom_forward(UpperCamelCase__			: Optional[int]    ):
        def custom_forward(*UpperCamelCase__			: int    ):
          return module(*UpperCamelCase__    )

        return custom_forward

      if is_torch_version(""">="""			,							"""1.11.0"""    ):
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block    )			,							UpperCamelCase__			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
        __magic_name__				     =       sample.to(UpperCamelCase__    )

        # up
        for up_block in self.up_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(
              create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
      else:
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block    )			,							UpperCamelCase__			,							UpperCamelCase__    )
        __magic_name__				     =       sample.to(UpperCamelCase__    )

        # up
        for up_block in self.up_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__			,							UpperCamelCase__    )
    else:
      # middle
      __magic_name__				     =       self.mid_block(UpperCamelCase__			,							UpperCamelCase__    )
      __magic_name__				     =       sample.to(UpperCamelCase__    )

      # up
      for up_block in self.up_blocks:
        __magic_name__				     =       up_block(UpperCamelCase__			,							UpperCamelCase__    )

        # post-process
    if latent_embeds is None:
      __magic_name__				     =       self.conv_norm_out(UpperCamelCase__    )
    else:
      __magic_name__				     =       self.conv_norm_out(UpperCamelCase__			,							UpperCamelCase__    )
    __magic_name__				     =       self.conv_act(UpperCamelCase__    )
    __magic_name__				     =       self.conv_out(UpperCamelCase__    )

    return sample







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: str			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Union[str, Any]=None			,							UpperCamelCase__			: Dict="random"			,							UpperCamelCase__			: List[Any]=False			,							UpperCamelCase__			: Dict=True    )			->	Optional[Any]:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       n_e
    __magic_name__				     =       vq_embed_dim
    __magic_name__				     =       beta
    __magic_name__				     =       legacy

    __magic_name__				     =       nn.Embedding(self.n_e			,							self.vq_embed_dim    )
    self.embedding.weight.data.uniform_(-1.0 / self.n_e			,							1.0 / self.n_e    )

    __magic_name__				     =       remap
    if self.remap is not None:
      self.register_buffer("""used"""			,							torch.tensor(np.load(self.remap    )    )    )
      __magic_name__				     =       self.used.shape[0]
      __magic_name__				     =       unknown_index  # "random" or "extra" or integer
      if self.unknown_index == "extra":
        __magic_name__				     =       self.re_embed
        __magic_name__				     =       self.re_embed + 1
      print(
          F'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
          F'''Using {self.unknown_index} for unknown indices.'''    )
    else:
      __magic_name__				     =       n_e

    __magic_name__				     =       sane_index_shape
  def 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: Dict    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       inds.shape
    assert len(UpperCamelCase__    ) > 1
    __magic_name__				     =       inds.reshape(ishape[0]			,							-1    )
    __magic_name__				     =       self.used.to(UpperCamelCase__    )
    __magic_name__				     =       (inds[:, :, None] == used[None, None, ...]).long()
    __magic_name__				     =       match.argmax(-1    )
    __magic_name__				     =       match.sum(2    ) < 1
    if self.unknown_index == "random":
      __magic_name__				     =       torch.randint(0			,							self.re_embed			,							size=new[unknown].shape    ).to(device=new.device    )
    else:
      __magic_name__				     =       self.unknown_index
    return new.reshape(UpperCamelCase__    )
  def 							_lowercase							( self			: Optional[Any]			,							UpperCamelCase__			: str    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       inds.shape
    assert len(UpperCamelCase__    ) > 1
    __magic_name__				     =       inds.reshape(ishape[0]			,							-1    )
    __magic_name__				     =       self.used.to(UpperCamelCase__    )
    if self.re_embed > self.used.shape[0]:  # extra token
      __magic_name__				     =       0  # simply set to zero
    __magic_name__				     =       torch.gather(used[None, :][inds.shape[0] * [0], :]			,							1			,							UpperCamelCase__    )
    return back.reshape(UpperCamelCase__    )
  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: List[str]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       z.permute(0			,							2			,							3			,							1    ).contiguous()
    __magic_name__				     =       z.view(-1			,							self.vq_embed_dim    )

    # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
    __magic_name__				     =       torch.argmin(torch.cdist(UpperCamelCase__			,							self.embedding.weight    )			,							dim=1    )

    __magic_name__				     =       self.embedding(UpperCamelCase__    ).view(z.shape    )
    __magic_name__				     =       None
    __magic_name__				     =       None

    # compute loss for embedding
    if not self.legacy:
      __magic_name__				     =       self.beta * torch.mean((z_q.detach() - z) ** 2    ) + torch.mean((z_q - z.detach()) ** 2    )
    else:
      __magic_name__				     =       torch.mean((z_q.detach() - z) ** 2    ) + self.beta * torch.mean((z_q - z.detach()) ** 2    )

    # preserve gradients
    __magic_name__				     =       z + (z_q - z).detach()

    # reshape back to match original input shape
    __magic_name__				     =       z_q.permute(0			,							3			,							1			,							2    ).contiguous()

    if self.remap is not None:
      __magic_name__				     =       min_encoding_indices.reshape(z.shape[0]			,							-1    )  # add batch axis
      __magic_name__				     =       self.remap_to_used(UpperCamelCase__    )
      __magic_name__				     =       min_encoding_indices.reshape(-1			,							1    )  # flatten

    if self.sane_index_shape:
      __magic_name__				     =       min_encoding_indices.reshape(z_q.shape[0]			,							z_q.shape[2]			,							z_q.shape[3]    )

    return z_q, loss, (perplexity, min_encodings, min_encoding_indices)





  def 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Union[str, Any]    )			->	int:







    """simple docstring"""






    if self.remap is not None:
      __magic_name__				     =       indices.reshape(shape[0]			,							-1    )  # add batch axis
      __magic_name__				     =       self.unmap_to_all(UpperCamelCase__    )
      __magic_name__				     =       indices.reshape(-1    )  # flatten again

    # get quantized latent vectors
    __magic_name__				     =       self.embedding(UpperCamelCase__    )

    if shape is not None:
      __magic_name__				     =       z_q.view(UpperCamelCase__    )
      # reshape back to match original input shape
      __magic_name__				     =       z_q.permute(0			,							3			,							1			,							2    ).contiguous()

    return z_q







class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  def __init__( self			: Dict			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Tuple=False    )			->	Optional[int]:







    """simple docstring"""






    __magic_name__				     =       parameters
    __magic_name__     ,    __magic_name__				     =       torch.chunk(UpperCamelCase__			,							2			,							dim=1    )
    __magic_name__				     =       torch.clamp(self.logvar			,							-30.0			,							20.0    )
    __magic_name__				     =       deterministic
    __magic_name__				     =       torch.exp(0.5 * self.logvar    )
    __magic_name__				     =       torch.exp(self.logvar    )
    if self.deterministic:
      __magic_name__				     =       __magic_name__				     =       torch.zeros_like(
          self.mean			,							device=self.parameters.device			,							dtype=self.parameters.dtype    )
  def 							_lowercase							( self			: Tuple			,							UpperCamelCase__			: Optional[torch.Generator] = None    )			->	torch.FloatTensor:







    """simple docstring"""






    __magic_name__				     =       randn_tensor(
        self.mean.shape			,							generator=UpperCamelCase__			,							device=self.parameters.device			,							dtype=self.parameters.dtype    )
    __magic_name__				     =       self.mean + self.std * sample
    return x
  def 							_lowercase							( self			: Dict			,							UpperCamelCase__			: Optional[int]=None    )			->	Any:







    """simple docstring"""






    if self.deterministic:
      return torch.Tensor([0.0]    )
    else:
      if other is None:
        return 0.5 * torch.sum(torch.pow(self.mean			,							2    ) + self.var - 1.0 - self.logvar			,							dim=[1, 2, 3]    )
      else:
        return 0.5 * torch.sum(
            torch.pow(self.mean - other.mean			,							2    ) / other.var
            + self.var / other.var
            - 1.0
            - self.logvar
            + other.logvar			,							dim=[1, 2, 3]			,							)
  def 							_lowercase							( self			: Optional[int]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Dict=[1, 2, 3]    )			->	Optional[int]:







    """simple docstring"""






    if self.deterministic:
      return torch.Tensor([0.0]    )
    __magic_name__				     =       np.log(2.0 * np.pi    )
    return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean			,							2    ) / self.var			,							dim=UpperCamelCase__    )





  def 							_lowercase							( self			: Optional[int]    )			->	Optional[int]:







    """simple docstring"""






    return self.mean

