

'''simple docstring'''
def  a__       (    _SCREAMING_SNAKE_CASE   :       list[int]	,				_SCREAMING_SNAKE_CASE   :       list[int]				)     ->					tuple[float, float]:


	"""simple docstring"""





	if not len(_SCREAMING_SNAKE_CASE				) == len(_SCREAMING_SNAKE_CASE				) == 3:
		raise ValueError("Please enter a valid equation."				)
	if equationa[0] == equationa[1] == equationa[0] == equationa[1] == 0:
		raise ValueError("Both a & b of two equations can't be zero."				)

	# Extract the coefficients
	UpperCAmelCase_					,					UpperCAmelCase_					,					UpperCAmelCase_     :      List[str]						=  equationa
	UpperCAmelCase_					,					UpperCAmelCase_					,					UpperCAmelCase_     :      List[Any]						=  equationa

	# Calculate the determinants of the matrices
	UpperCAmelCase_     :      Optional[int]						=  aa * ba - aa * ba
	UpperCAmelCase_     :      Optional[int]						=  ca * ba - ca * ba
	UpperCAmelCase_     :      Any						=  aa * ca - aa * ca

	# Check if the system of linear equations has a solution (using Cramer's rule)
	if determinant == 0:
		if determinant_x == determinant_y == 0:
			raise ValueError("Infinite solutions. (Consistent system)"				)
		else:
			raise ValueError("No solution. (Inconsistent system)"				)
	else:
		if determinant_x == determinant_y == 0:
			# Trivial solution (Inconsistent system)
			return (0.0, 0.0)
		else:
			UpperCAmelCase_     :      Optional[int]						=  determinant_x / determinant
			UpperCAmelCase_     :      List[Any]						=  determinant_y / determinant
			# Non-Trivial Solution (Consistent system)
			return (x, y)








'''simple docstring'''
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 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       torch.FloatTensor






class 					__lowerCAmelCase   (   nn.Module						):

						"""simple docstring"""






						def __init__(  self							:      str ,  lowerCAmelCase							:      Tuple=3 ,  lowerCAmelCase							:      Union[str, Any]=3 ,  lowerCAmelCase							:      str=("DownEncoderBlock2D",) ,  lowerCAmelCase							:      Union[str, Any]=(64,) ,  lowerCAmelCase							:      str=2 ,  lowerCAmelCase							:      List[Any]=32 ,  lowerCAmelCase							:      List[str]="silu" ,  lowerCAmelCase							:      Tuple=True ,  ):
													super().__init__()
													A_			    =		layers_per_block

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

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

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

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

													# mid
													A_			    =		UNetMidBlockaD(
													    in_channels=block_out_channels[-1] ,  resnet_eps=1e-6 ,  resnet_act_fn=lowerCAmelCase ,  output_scale_factor=1 ,  resnet_time_scale_shift="default" ,  attention_head_dim=block_out_channels[-1] ,  resnet_groups=lowerCAmelCase ,  temb_channels=lowerCAmelCase ,  )

													# out
													A_			    =		nn.GroupNorm(num_channels=block_out_channels[-1] ,  num_groups=lowerCAmelCase ,  eps=1e-6	)
													A_			    =		nn.SiLU()

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

													A_			    =		False






						def 				_UpperCAmelCase       (  self							:      Dict ,  lowerCAmelCase							:      int	):
													A_			    =		x
													A_			    =		self.conv_in(lowerCAmelCase	)

													if self.training and self.gradient_checkpointing:

																				def create_custom_forward(lowerCAmelCase							:      Dict	):
																											def custom_forward(*lowerCAmelCase							:      List[Any]	):
																																		return module(*lowerCAmelCase	)

																											return custom_forward

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

													else:
																				# down
																				for down_block in self.down_blocks:
																											A_			    =		down_block(lowerCAmelCase	)

																				# middle
																				A_			    =		self.mid_block(lowerCAmelCase	)

													# post-process
													A_			    =		self.conv_norm_out(lowerCAmelCase	)
													A_			    =		self.conv_act(lowerCAmelCase	)
													A_			    =		self.conv_out(lowerCAmelCase	)

													return sample






class 					__lowerCAmelCase   (   nn.Module						):

						"""simple docstring"""






						def __init__(  self							:      Optional[int] ,  lowerCAmelCase							:      Optional[Any]=3 ,  lowerCAmelCase							:      Dict=3 ,  lowerCAmelCase							:      Dict=("UpDecoderBlock2D",) ,  lowerCAmelCase							:      Optional[int]=(64,) ,  lowerCAmelCase							:      str=2 ,  lowerCAmelCase							:      Dict=32 ,  lowerCAmelCase							:      Any="silu" ,  lowerCAmelCase							:      List[Any]="group" ,  ):
													super().__init__()
													A_			    =		layers_per_block

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

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

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

													# mid
													A_			    =		UNetMidBlockaD(
													    in_channels=block_out_channels[-1] ,  resnet_eps=1e-6 ,  resnet_act_fn=lowerCAmelCase ,  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=lowerCAmelCase ,  temb_channels=lowerCAmelCase ,  )

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

																				A_			    =		i == len(lowerCAmelCase	) - 1

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

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

													A_			    =		False






						def 				_UpperCAmelCase       (  self							:      str ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      Union[str, Any]=None	):
													A_			    =		z
													A_			    =		self.conv_in(lowerCAmelCase	)

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

																				def create_custom_forward(lowerCAmelCase							:      Union[str, Any]	):
																											def custom_forward(*lowerCAmelCase							:      Union[str, Any]	):
																																		return module(*lowerCAmelCase	)

																											return custom_forward

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

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

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

																				# up
																				for up_block in self.up_blocks:
																											A_			    =		up_block(lowerCAmelCase ,  lowerCAmelCase	)

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

													return sample






class 					__lowerCAmelCase   (   nn.Module						):

						"""simple docstring"""






						def __init__(  self							:      Dict ,  lowerCAmelCase							:      Dict ,  lowerCAmelCase							:      Union[str, Any] ,  lowerCAmelCase							:      Optional[Any] ,  lowerCAmelCase							:      List[str]=None ,  lowerCAmelCase							:      str="random" ,  lowerCAmelCase							:      Union[str, Any]=False ,  lowerCAmelCase							:      Optional[Any]=True	):
													super().__init__()
													A_			    =		n_e
													A_			    =		vq_embed_dim
													A_			    =		beta
													A_			    =		legacy

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

													A_			    =		remap
													if self.remap is not None:
																				self.register_buffer("used" ,  torch.tensor(np.load(self.remap	)	)	)
																				A_			    =		self.used.shape[0]
																				A_			    =		unknown_index  # "random" or "extra" or integer
																				if self.unknown_index == "extra":
																											A_			    =		self.re_embed
																											A_			    =		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:
																				A_			    =		n_e

													A_			    =		sane_index_shape






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






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






						def 				_UpperCAmelCase       (  self							:      int ,  lowerCAmelCase							:      List[Any]	):
													# reshape z -> (batch, height, width, channel) and flatten
													A_			    =		z.permute(0 ,  2 ,  3 ,  1	).contiguous()
													A_			    =		z.view(-1 ,  self.vq_embed_dim	)

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

													A_			    =		self.embedding(lowerCAmelCase	).view(z.shape	)
													A_			    =		None
													A_			    =		None

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

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

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

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

													if self.sane_index_shape:
																				A_			    =		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 				_UpperCAmelCase       (  self							:      str ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      int	):
													# shape specifying (batch, height, width, channel)
													if self.remap is not None:
																				A_			    =		indices.reshape(shape[0] ,  -1	)  # add batch axis
																				A_			    =		self.unmap_to_all(lowerCAmelCase	)
																				A_			    =		indices.reshape(-1	)  # flatten again

													# get quantized latent vectors
													A_			    =		self.embedding(lowerCAmelCase	)

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

													return z_q






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						def __init__(  self							:      Dict ,  lowerCAmelCase							:      Union[str, Any] ,  lowerCAmelCase							:      Union[str, Any]=False	):
													A_			    =		parameters
													A_    ,					A_			    =		torch.chunk(lowerCAmelCase ,  2 ,  dim=1	)
													A_			    =		torch.clamp(self.logvar ,  -3_0.0 ,  2_0.0	)
													A_			    =		deterministic
													A_			    =		torch.exp(0.5 * self.logvar	)
													A_			    =		torch.exp(self.logvar	)
													if self.deterministic:
																				A_			    =		A_			    =		torch.zeros_like(
																				    self.mean ,  device=self.parameters.device ,  dtype=self.parameters.dtype	)






						def 				_UpperCAmelCase       (  self							:      List[str] ,  lowerCAmelCase							:      Optional[torch.Generator] = None	):
													# make sure sample is on the same device as the parameters and has same dtype
													A_			    =		randn_tensor(
													    self.mean.shape ,  generator=lowerCAmelCase ,  device=self.parameters.device ,  dtype=self.parameters.dtype	)
													A_			    =		self.mean + self.std * sample
													return x






						def 				_UpperCAmelCase       (  self							:      Optional[int] ,  lowerCAmelCase							:      str=None	):
													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 				_UpperCAmelCase       (  self							:      List[Any] ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      Optional[int]=[1, 2, 3]	):
													if self.deterministic:
																				return torch.Tensor([0.0]	)
													A_			    =		np.log(2.0 * np.pi	)
													return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean ,  2	) / self.var ,  dim=lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Union[str, Any]	):
													return self.mean





