




"""simple docstring"""

import types
from typing import List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
from transformers.models.clip.modeling_clip import CLIPTextModelOutput

from diffusers.models import PriorTransformer
from diffusers.pipelines import DiffusionPipeline, StableDiffusionImageVariationPipeline
from diffusers.schedulers import UnCLIPScheduler
from diffusers.utils import logging, randn_tensor


__lowerCamelCase								=						logging.get_logger(__name__)  # pylint: disable=invalid-name



def 				a (						self     :				Any,    __snake_case     :				Any,    __snake_case     :				Tuple,    __snake_case     :				Optional[int],    __snake_case     :				List[Any]  ):



    '''simple docstring'''





    UpperCAmelCase_					:Dict      			=      image.to(device=_UpperCAmelCase  )
    UpperCAmelCase_					:List[Any]      			=      image  # take image as image_embeddings
    UpperCAmelCase_					:Dict      			=      image_embeddings.unsqueeze(1  )

    # duplicate image embeddings for each generation per prompt, using mps friendly method
    UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_					:Any      			=      image_embeddings.shape
    UpperCAmelCase_					:int      			=      image_embeddings.repeat(1,    _UpperCAmelCase,    1  )
    UpperCAmelCase_					:List[str]      			=      image_embeddings.view(bs_embed * num_images_per_prompt,    _UpperCAmelCase,    -1  )

    if do_classifier_free_guidance:
        UpperCAmelCase_					:int      			=      torch.zeros_like(_UpperCAmelCase  )

        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        UpperCAmelCase_					:Optional[Any]      			=      torch.cat([uncond_embeddings, image_embeddings]  )

    return image_embeddings





class  _snake_case				(							lowerCamelCase__       ):

    '''simple docstring'''




    def __init__(							self      :  Union[str, Any]					,							snake_case      :  PriorTransformer					,							snake_case      :  CLIPTokenizer					,							snake_case      :  CLIPTextModelWithProjection					,							snake_case      :  UnCLIPScheduler					,							snake_case      :  Optional[dict] = None					,							):
        super().__init__()

        UpperCAmelCase_					:Optional[int]      			=      {'''image_encoder''': None} if decoder_pipe_kwargs is None else decoder_pipe_kwargs

        UpperCAmelCase_					:List[Any]      			=      decoder_pipe_kwargs.get('''torch_dtype'''					,							SCREAMING_SNAKE_CASE_     ) or prior.dtype

        UpperCAmelCase_					:Union[str, Any]      			=      StableDiffusionImageVariationPipeline.from_pretrained(
            '''lambdalabs/sd-image-variations-diffusers'''					,							**SCREAMING_SNAKE_CASE_     )

        # replace `_encode_image` method
        UpperCAmelCase_					:List[str]      			=      types.MethodType(_encode_image					,							self.decoder_pipe     )

        self.register_modules(
            prior=SCREAMING_SNAKE_CASE_					,							tokenizer=SCREAMING_SNAKE_CASE_					,							text_encoder=SCREAMING_SNAKE_CASE_					,							prior_scheduler=SCREAMING_SNAKE_CASE_					,							)





    def 				snake_case_ (							self      :  List[str]					,							snake_case      :  List[str]					,							snake_case      :  str					,							snake_case      :  List[str]					,							snake_case      :  Union[str, Any]					,							snake_case      :  Optional[Union[CLIPTextModelOutput, Tuple]] = None					,							snake_case      :  Optional[torch.Tensor] = None					,							):
        if text_model_output is None:
            UpperCAmelCase_					:Any      			=      len(SCREAMING_SNAKE_CASE_     ) if isinstance(SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_     ) else 1
            # get prompt text embeddings
            UpperCAmelCase_					:Dict      			=      self.tokenizer(
                SCREAMING_SNAKE_CASE_					,							padding='''max_length'''					,							max_length=self.tokenizer.model_max_length					,							return_tensors='''pt'''					,							)
            UpperCAmelCase_					:Any      			=      text_inputs.input_ids
            UpperCAmelCase_					:str      			=      text_inputs.attention_mask.bool().to(SCREAMING_SNAKE_CASE_     )

            if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
                UpperCAmelCase_					:Any      			=      self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]     )
                logger.warning(
                    '''The following part of your input was truncated because CLIP can only handle sequences up to'''
                    f' {self.tokenizer.model_max_length} tokens: {removed_text}'     )
                UpperCAmelCase_					:List[Any]      			=      text_input_ids[:, : self.tokenizer.model_max_length]

            UpperCAmelCase_					:Dict      			=      self.text_encoder(text_input_ids.to(SCREAMING_SNAKE_CASE_     )     )

            UpperCAmelCase_					:Optional[Any]      			=      text_encoder_output.text_embeds
            UpperCAmelCase_					:List[str]      			=      text_encoder_output.last_hidden_state

        else:
            UpperCAmelCase_					:Any      			=      text_model_output[0].shape[0]
            UpperCAmelCase_			,UpperCAmelCase_					:Dict      			=      text_model_output[0], text_model_output[1]
            UpperCAmelCase_					:List[str]      			=      text_attention_mask

        UpperCAmelCase_					:Dict      			=      text_embeddings.repeat_interleave(SCREAMING_SNAKE_CASE_					,							dim=0     )
        UpperCAmelCase_					:str      			=      text_encoder_hidden_states.repeat_interleave(SCREAMING_SNAKE_CASE_					,							dim=0     )
        UpperCAmelCase_					:Tuple      			=      text_mask.repeat_interleave(SCREAMING_SNAKE_CASE_					,							dim=0     )

        if do_classifier_free_guidance:
            UpperCAmelCase_					:Dict      			=      [''''''] * batch_size

            UpperCAmelCase_					:Union[str, Any]      			=      self.tokenizer(
                SCREAMING_SNAKE_CASE_					,							padding='''max_length'''					,							max_length=self.tokenizer.model_max_length					,							truncation=SCREAMING_SNAKE_CASE_					,							return_tensors='''pt'''					,							)
            UpperCAmelCase_					:List[Any]      			=      uncond_input.attention_mask.bool().to(SCREAMING_SNAKE_CASE_     )
            UpperCAmelCase_					:str      			=      self.text_encoder(uncond_input.input_ids.to(SCREAMING_SNAKE_CASE_     )     )

            UpperCAmelCase_					:Union[str, Any]      			=      uncond_embeddings_text_encoder_output.text_embeds
            UpperCAmelCase_					:Optional[int]      			=      uncond_embeddings_text_encoder_output.last_hidden_state

            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

            UpperCAmelCase_					:List[Any]      			=      uncond_embeddings.shape[1]
            UpperCAmelCase_					:Optional[int]      			=      uncond_embeddings.repeat(1					,							SCREAMING_SNAKE_CASE_     )
            UpperCAmelCase_					:str      			=      uncond_embeddings.view(batch_size * num_images_per_prompt					,							SCREAMING_SNAKE_CASE_     )

            UpperCAmelCase_					:Optional[int]      			=      uncond_text_encoder_hidden_states.shape[1]
            UpperCAmelCase_					:Optional[Any]      			=      uncond_text_encoder_hidden_states.repeat(1					,							SCREAMING_SNAKE_CASE_					,							1     )
            UpperCAmelCase_					:str      			=      uncond_text_encoder_hidden_states.view(
                batch_size * num_images_per_prompt					,							SCREAMING_SNAKE_CASE_					,							-1     )
            UpperCAmelCase_					:str      			=      uncond_text_mask.repeat_interleave(SCREAMING_SNAKE_CASE_					,							dim=0     )

            # done duplicates

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and text embeddings into a single batch
            # to avoid doing two forward passes
            UpperCAmelCase_					:Dict      			=      torch.cat([uncond_embeddings, text_embeddings]     )
            UpperCAmelCase_					:Optional[Any]      			=      torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]     )

            UpperCAmelCase_					:str      			=      torch.cat([uncond_text_mask, text_mask]     )

        return text_embeddings, text_encoder_hidden_states, text_mask





    @property
    def 				snake_case_ (							self      :  Optional[Any]     ):
        if self.device != torch.device('''meta'''     ) or not hasattr(self.prior					,							'''_hf_hook'''     ):
            return self.device
        for module in self.prior.modules():
            if (
                hasattr(SCREAMING_SNAKE_CASE_					,							'''_hf_hook'''     )
                and hasattr(module._hf_hook					,							'''execution_device'''     )
                and module._hf_hook.execution_device is not None
            ):
                return torch.device(module._hf_hook.execution_device     )
        return self.device





    def 				snake_case_ (							self      :  Union[str, Any]					,							snake_case      :  Dict					,							snake_case      :  int					,							snake_case      :  Optional[int]					,							snake_case      :  str					,							snake_case      :  str					,							snake_case      :  Optional[int]     ):
        if latents is None:
            UpperCAmelCase_					:int      			=      randn_tensor(SCREAMING_SNAKE_CASE_					,							generator=SCREAMING_SNAKE_CASE_					,							device=SCREAMING_SNAKE_CASE_					,							dtype=SCREAMING_SNAKE_CASE_     )
        else:
            if latents.shape != shape:
                raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'     )
            UpperCAmelCase_					:int      			=      latents.to(SCREAMING_SNAKE_CASE_     )

        UpperCAmelCase_					:List[Any]      			=      latents * scheduler.init_noise_sigma
        return latents





    def 				snake_case_ (							self      :  List[Any]					,							snake_case      :  Optional[Union[str, torch.device]] = None     ):
        self.decoder_pipe.to(SCREAMING_SNAKE_CASE_     )
        super().to(SCREAMING_SNAKE_CASE_     )





    @torch.no_grad()
    def __call__(							self      :  Any					,							snake_case      :  Optional[Union[str, List[str]]] = None					,							snake_case      :  Optional[int] = None					,							snake_case      :  Optional[int] = None					,							snake_case      :  int = 1					,							snake_case      :  int = 25					,							snake_case      :  Optional[torch.Generator] = None					,							snake_case      :  Optional[torch.FloatTensor] = None					,							snake_case      :  Optional[Union[CLIPTextModelOutput, Tuple]] = None					,							snake_case      :  Optional[torch.Tensor] = None					,							snake_case      :  float = 4.0					,							snake_case      :  float = 8.0					,							snake_case      :  int = 50					,							snake_case      :  Optional[int] = 1					,							snake_case      :  float = 0.0					,							snake_case      :  Optional[str] = "pil"					,							snake_case      :  bool = True					,							):
        if prompt is not None:
            if isinstance(SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_     ):
                UpperCAmelCase_					:Union[str, Any]      			=      1
            elif isinstance(SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_     ):
                UpperCAmelCase_					:List[Any]      			=      len(SCREAMING_SNAKE_CASE_     )
            else:
                raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(SCREAMING_SNAKE_CASE_     )}'     )
        else:
            UpperCAmelCase_					:str      			=      text_model_output[0].shape[0]

        UpperCAmelCase_					:Tuple      			=      self._execution_device

        UpperCAmelCase_					:Optional[Any]      			=      batch_size * num_images_per_prompt

        UpperCAmelCase_					:Dict      			=      prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0

        UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_					:List[str]      			=      self._encode_prompt(
            SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_     )

        # prior

        self.prior_scheduler.set_timesteps(SCREAMING_SNAKE_CASE_					,							device=SCREAMING_SNAKE_CASE_     )
        UpperCAmelCase_					:Any      			=      self.prior_scheduler.timesteps

        UpperCAmelCase_					:int      			=      self.prior.config.embedding_dim

        UpperCAmelCase_					:Optional[int]      			=      self.prepare_latents(
            (batch_size, embedding_dim)					,							text_embeddings.dtype					,							SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_					,							self.prior_scheduler					,							)

        for i, t in enumerate(self.progress_bar(SCREAMING_SNAKE_CASE_     )     ):
            # expand the latents if we are doing classifier free guidance
            UpperCAmelCase_					:Any      			=      torch.cat([prior_latents] * 2     ) if do_classifier_free_guidance else prior_latents

            UpperCAmelCase_					:str      			=      self.prior(
                SCREAMING_SNAKE_CASE_					,							timestep=SCREAMING_SNAKE_CASE_					,							proj_embedding=SCREAMING_SNAKE_CASE_					,							encoder_hidden_states=SCREAMING_SNAKE_CASE_					,							attention_mask=SCREAMING_SNAKE_CASE_					,							).predicted_image_embedding

            if do_classifier_free_guidance:
                UpperCAmelCase_			,UpperCAmelCase_					:Optional[Any]      			=      predicted_image_embedding.chunk(2     )
                UpperCAmelCase_					:List[str]      			=      predicted_image_embedding_uncond + prior_guidance_scale * (
                    predicted_image_embedding_text - predicted_image_embedding_uncond
                )

            if i + 1 == prior_timesteps_tensor.shape[0]:
                UpperCAmelCase_					:Any      			=      None
            else:
                UpperCAmelCase_					:Optional[int]      			=      prior_timesteps_tensor[i + 1]

            UpperCAmelCase_					:Any      			=      self.prior_scheduler.step(
                SCREAMING_SNAKE_CASE_					,							timestep=SCREAMING_SNAKE_CASE_					,							sample=SCREAMING_SNAKE_CASE_					,							generator=SCREAMING_SNAKE_CASE_					,							prev_timestep=SCREAMING_SNAKE_CASE_					,							).prev_sample

        UpperCAmelCase_					:Union[str, Any]      			=      self.prior.post_process_latents(SCREAMING_SNAKE_CASE_     )

        UpperCAmelCase_					:str      			=      prior_latents

        UpperCAmelCase_					:Dict      			=      self.decoder_pipe(
            image=SCREAMING_SNAKE_CASE_					,							height=SCREAMING_SNAKE_CASE_					,							width=SCREAMING_SNAKE_CASE_					,							num_inference_steps=SCREAMING_SNAKE_CASE_					,							guidance_scale=SCREAMING_SNAKE_CASE_					,							generator=SCREAMING_SNAKE_CASE_					,							output_type=SCREAMING_SNAKE_CASE_					,							return_dict=SCREAMING_SNAKE_CASE_					,							num_images_per_prompt=SCREAMING_SNAKE_CASE_					,							eta=SCREAMING_SNAKE_CASE_					,							)
        return output








"""simple docstring"""

from typing import Dict, Tuple, overload

import torch
import torch.types
from torch import nn

from . import residue_constants as rc
from .rigid_utils import Rigid, Rotation
from .tensor_utils import batched_gather



@overload
def 				a (						__snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				None  ):



				'''simple docstring'''





				...



@overload
def 				a (						__snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor  ):



				'''simple docstring'''





				...



def 				a (						__snake_case     :				List[str],    __snake_case     :				Optional[int],    __snake_case     :				str  ):



				'''simple docstring'''





				UpperCAmelCase_					:int      			=      aatype == rc.restype_order['''G''']
				UpperCAmelCase_					:Optional[Any]      			=      rc.atom_order['''CA''']
				UpperCAmelCase_					:Optional[Any]      			=      rc.atom_order['''CB''']
				UpperCAmelCase_					:Optional[Any]      			=      torch.where(
				    is_gly[..., None].expand(*((-1,) * len(is_gly.shape  )),    3  ),    all_atom_positions[..., ca_idx, :],    all_atom_positions[..., cb_idx, :],    )

				if all_atom_masks is not None:
								UpperCAmelCase_					:Any      			=      torch.where(
								    __snake_case,    all_atom_masks[..., ca_idx],    all_atom_masks[..., cb_idx],    )
								return pseudo_beta, pseudo_beta_mask
				else:
								return pseudo_beta



def 				a (						__snake_case     :				torch.Tensor,    __snake_case     :				Dict[str, torch.Tensor]  ):



				'''simple docstring'''





				UpperCAmelCase_					:List[Any]      			=      batched_gather(
				    __snake_case,    batch['''residx_atom37_to_atom14'''],    dim=-2,    no_batch_dims=len(atomaa.shape[:-2]  ),    )

				UpperCAmelCase_					:Optional[Any]      			=      atomaa_data * batch['''atom37_atom_exists'''][..., None]

				return atomaa_data



def 				a (						__snake_case     :				Dict[str, torch.Tensor]  ):



				'''simple docstring'''





				UpperCAmelCase_					:str      			=      template_feats['''template_aatype''']
				UpperCAmelCase_					:Union[str, Any]      			=      template_feats['''template_torsion_angles_sin_cos''']
				UpperCAmelCase_					:Tuple      			=      template_feats['''template_alt_torsion_angles_sin_cos''']
				UpperCAmelCase_					:Tuple      			=      template_feats['''template_torsion_angles_mask''']
				UpperCAmelCase_					:int      			=      torch.cat(
				    [
				        nn.functional.one_hot(__snake_case,    22  ),
				        torsion_angles_sin_cos.reshape(*torsion_angles_sin_cos.shape[:-2],    14  ),
				        alt_torsion_angles_sin_cos.reshape(*alt_torsion_angles_sin_cos.shape[:-2],    14  ),
				        torsion_angles_mask,
				    ],    dim=-1,    )

				return template_angle_feat



def 				a (						__snake_case     :				Dict[str, torch.Tensor],    __snake_case     :				torch.types.Number,    __snake_case     :				torch.types.Number,    __snake_case     :				int,    __snake_case     :				bool = False,    __snake_case     :				float = 1E-20,    __snake_case     :				float = 1E8,    ):



				'''simple docstring'''





				UpperCAmelCase_					:List[str]      			=      batch['''template_pseudo_beta_mask''']
				UpperCAmelCase_					:Any      			=      template_mask[..., None] * template_mask[..., None, :]

				# Compute distogram (this seems to differ slightly from Alg. 5)
				UpperCAmelCase_					:List[str]      			=      batch['''template_pseudo_beta''']
				UpperCAmelCase_					:Any      			=      torch.sum((tpb[..., None, :] - tpb[..., None, :, :]) ** 2,    dim=-1,    keepdim=__snake_case  )
				UpperCAmelCase_					:Any      			=      torch.linspace(__snake_case,    __snake_case,    __snake_case,    device=tpb.device  ) ** 2
				UpperCAmelCase_					:Union[str, Any]      			=      torch.cat([lower[1:], lower.new_tensor([inf]  )],    dim=-1  )
				UpperCAmelCase_					:Tuple      			=      ((dgram > lower) * (dgram < upper)).type(dgram.dtype  )

				UpperCAmelCase_					:List[str]      			=      [dgram, template_mask_ad[..., None]]

				UpperCAmelCase_					:torch.LongTensor      			=      nn.functional.one_hot(
				    batch['''template_aatype'''],    rc.restype_num + 2,    )

				UpperCAmelCase_					:Optional[Any]      			=      batch['''template_aatype'''].shape[-1]
				to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2],    __snake_case,    -1,    -1  )  )
				to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2],    -1,    __snake_case,    -1  )  )

				UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_					:Optional[int]      			=      [rc.atom_order[a] for a in ['''N''', '''CA''', '''C''']]
				UpperCAmelCase_					:str      			=      Rigid.make_transform_from_reference(
				    n_xyz=batch['''template_all_atom_positions'''][..., n, :],    ca_xyz=batch['''template_all_atom_positions'''][..., ca, :],    c_xyz=batch['''template_all_atom_positions'''][..., c, :],    eps=__snake_case,    )
				UpperCAmelCase_					:Dict      			=      rigids.get_trans()[..., None, :, :]
				UpperCAmelCase_					:Dict      			=      rigids[..., None].invert_apply(__snake_case  )

				UpperCAmelCase_					:Optional[int]      			=      torch.rsqrt(eps + torch.sum(rigid_vec**2,    dim=-1  )  )

				UpperCAmelCase_					:Optional[int]      			=      batch['''template_all_atom_mask''']
				UpperCAmelCase_					:Tuple      			=      t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
				UpperCAmelCase_					:int      			=      template_mask[..., None] * template_mask[..., None, :]

				UpperCAmelCase_					:str      			=      inv_distance_scalar * template_mask_ad
				UpperCAmelCase_					:List[Any]      			=      rigid_vec * inv_distance_scalar[..., None]

				if not use_unit_vector:
								UpperCAmelCase_					:List[Any]      			=      unit_vector * 0.0

				to_concat.extend(torch.unbind(unit_vector[..., None, :],    dim=-1  )  )
				to_concat.append(template_mask_ad[..., None]  )

				UpperCAmelCase_					:Optional[int]      			=      torch.cat(__snake_case,    dim=-1  )
				UpperCAmelCase_					:Optional[Any]      			=      act * template_mask_ad[..., None]

				return act



def 				a (						__snake_case     :				Dict[str, torch.Tensor]  ):



				'''simple docstring'''





				UpperCAmelCase_					:torch.LongTensor      			=      nn.functional.one_hot(batch['''extra_msa'''],    23  )
				UpperCAmelCase_					:Union[str, Any]      			=      [
				    msa_ahot,
				    batch['''extra_has_deletion'''].unsqueeze(-1  ),
				    batch['''extra_deletion_value'''].unsqueeze(-1  ),
				]
				return torch.cat(__snake_case,    dim=-1  )



def 				a (						__snake_case     :				Rigid,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    ):



				'''simple docstring'''





				UpperCAmelCase_					:Optional[int]      			=      rrgdf[aatype, ...]

				# [*, N, 8] transformations, i.e.
				#   One [*, N, 8, 3, 3] rotation matrix and
				#   One [*, N, 8, 3]    translation matrix
				UpperCAmelCase_					:Optional[int]      			=      r.from_tensor_axa(__snake_case  )

				UpperCAmelCase_					:List[Any]      			=      alpha.new_zeros((*((1,) * len(alpha.shape[:-1]  )), 2)  )
				UpperCAmelCase_					:List[Any]      			=      1

				# [*, N, 8, 2]
				UpperCAmelCase_					:Tuple      			=      torch.cat([bb_rot.expand(*alpha.shape[:-2],    -1,    -1  ), alpha],    dim=-2  )

				# [*, N, 8, 3, 3]
				# Produces rotation matrices of the form:
				# [
				#   [1, 0  , 0  ],
				#   [0, a_2,-a_1],
				#   [0, a_1, a_2]
				# ]
				# This follows the original code rather than the supplement, which uses
				# different indices.

				UpperCAmelCase_					:Any      			=      alpha.new_zeros(default_r.get_rots().get_rot_mats().shape  )
				UpperCAmelCase_					:Optional[Any]      			=      1
				UpperCAmelCase_					:int      			=      alpha[..., 1]
				UpperCAmelCase_					:Optional[int]      			=      -alpha[..., 0]
				UpperCAmelCase_					:int      			=      alpha

				UpperCAmelCase_					:List[str]      			=      default_r.compose(Rigid(Rotation(rot_mats=__snake_case  ),    __snake_case  )  )

				UpperCAmelCase_					:List[str]      			=      all_frames[..., 5]
				UpperCAmelCase_					:str      			=      all_frames[..., 6]
				UpperCAmelCase_					:List[Any]      			=      all_frames[..., 7]

				UpperCAmelCase_					:Tuple      			=      all_frames[..., 4]
				UpperCAmelCase_					:Optional[Any]      			=      chia_frame_to_bb.compose(__snake_case  )
				UpperCAmelCase_					:Optional[Any]      			=      chia_frame_to_bb.compose(__snake_case  )
				UpperCAmelCase_					:Union[str, Any]      			=      chia_frame_to_bb.compose(__snake_case  )

				UpperCAmelCase_					:List[str]      			=      Rigid.cat(
				    [
				        all_frames[..., :5],
				        chia_frame_to_bb.unsqueeze(-1  ),
				        chia_frame_to_bb.unsqueeze(-1  ),
				        chia_frame_to_bb.unsqueeze(-1  ),
				    ],    dim=-1,    )

				UpperCAmelCase_					:Any      			=      r[..., None].compose(__snake_case  )

				return all_frames_to_global





def 				a (						__snake_case     :				Rigid,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    ):



				'''simple docstring'''





				UpperCAmelCase_					:Optional[Any]      			=      group_idx[aatype, ...]

				# [*, N, 14, 8]
				UpperCAmelCase_					:torch.LongTensor      			=      nn.functional.one_hot(
				    __snake_case,    num_classes=default_frames.shape[-3],    )

				# [*, N, 14, 8]
				UpperCAmelCase_					:Dict      			=      r[..., None, :] * group_mask_one_hot

				# [*, N, 14]
				UpperCAmelCase_					:List[Any]      			=      t_atoms_to_global.map_tensor_fn(lambda __snake_case  : torch.sum(__snake_case,    dim=-1  )  )

				# [*, N, 14, 1]
				UpperCAmelCase_					:List[Any]      			=      atom_mask[aatype, ...].unsqueeze(-1  )

				# [*, N, 14, 3]
				UpperCAmelCase_					:List[Any]      			=      lit_positions[aatype, ...]
				UpperCAmelCase_					:Optional[int]      			=      t_atoms_to_global.apply(__snake_case  )
				UpperCAmelCase_					:int      			=      pred_positions * atom_mask

				return pred_positions


