




"""simple docstring"""

import os
import pickle
import tempfile
import unittest
from typing import Callable, Optional

import numpy as np

from transformers import (
    BatchEncoding,
    BertTokenizer,
    BertTokenizerFast,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    TensorType,
    TokenSpan,
    is_tokenizers_available,
)
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import CaptureStderr, require_flax, require_tf, require_tokenizers, require_torch, slow


if is_tokenizers_available():
      from tokenizers import Tokenizer
      from tokenizers.models import WordPiece



class  _snake_case				(							unittest.TestCase       ):

    '''simple docstring'''




    def 				snake_case_ (							self      :  Any					,							snake_case      :  Any     ):
        UpperCAmelCase_					:str      			=      list(tokenizer_class.max_model_input_sizes.keys()     )
        for model_name in sa_models[:1]:
            UpperCAmelCase_					:Tuple      			=      tokenizer_class.from_pretrained(snake_case     )
            self.assertIsNotNone(snake_case     )
            self.assertIsInstance(snake_case					,							snake_case     )
            self.assertIsInstance(snake_case					,							snake_case     )

            for special_tok in tokenizer.all_special_tokens:
                self.assertIsInstance(snake_case					,							snake_case     )
                UpperCAmelCase_					:Optional[Any]      			=      tokenizer.convert_tokens_to_ids(snake_case     )
                self.assertIsInstance(snake_case					,							snake_case     )





    def 				snake_case_ (							self      :  List[Any]					,							snake_case      :  BatchEncoding					,							snake_case      :  Optional[Callable] = None     ):
        UpperCAmelCase_					:Any      			=      pickle.dumps(snake_case     )
        self.assertIsNotNone(snake_case     )

        UpperCAmelCase_					:List[Any]      			=      pickle.loads(snake_case     )

        # Ensure is_fast is correctly restored
        self.assertEqual(be_restored.is_fast					,							be_original.is_fast     )

        # Ensure encodings are potentially correctly restored
        if be_original.is_fast:
            self.assertIsNotNone(be_restored.encodings     )
        else:
            self.assertIsNone(be_restored.encodings     )

        # Ensure the keys are the same
        for original_v, restored_v in zip(be_original.values()					,							be_restored.values()     ):
            if equal_op:
                self.assertTrue(equal_op(snake_case					,							snake_case     )     )
            else:
                self.assertEqual(snake_case					,							snake_case     )





    @slow
    def 				snake_case_ (							self      :  List[Any]     ):
        self.check_tokenizer_from_pretrained(snake_case     )





    def 				snake_case_ (							self      :  Dict     ):
        self.assertEqual(TensorType('''tf'''     )					,							TensorType.TENSORFLOW     )
        self.assertEqual(TensorType('''pt'''     )					,							TensorType.PYTORCH     )
        self.assertEqual(TensorType('''np'''     )					,							TensorType.NUMPY     )





    @require_tokenizers
    def 				snake_case_ (							self      :  List[Any]     ):
        import numpy as np

        UpperCAmelCase_					:Tuple      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )
        UpperCAmelCase_					:Union[str, Any]      			=      BertTokenizerFast.from_pretrained('''bert-base-cased'''     )

        # Python no tensor
        with self.subTest('''BatchEncoding (Python, return_tensors=None)'''     ):
            self.assert_dump_and_restore(tokenizer_p('''Small example to encode'''     )     )

        with self.subTest('''BatchEncoding (Python, return_tensors=NUMPY)'''     ):
            self.assert_dump_and_restore(
                tokenizer_p('''Small example to encode'''					,							return_tensors=TensorType.NUMPY     )					,							np.array_equal     )

        with self.subTest('''BatchEncoding (Rust, return_tensors=None)'''     ):
            self.assert_dump_and_restore(tokenizer_r('''Small example to encode'''     )     )

        with self.subTest('''BatchEncoding (Rust, return_tensors=NUMPY)'''     ):
            self.assert_dump_and_restore(
                tokenizer_r('''Small example to encode'''					,							return_tensors=TensorType.NUMPY     )					,							np.array_equal     )





    @require_tf
    @require_tokenizers
    def 				snake_case_ (							self      :  Optional[int]     ):
        import tensorflow as tf

        def tf_array_equals(snake_case      :  str					,							snake_case      :  List[Any]     ):
            return tf.reduce_all(tf.equal(snake_case					,							snake_case     )     )

        UpperCAmelCase_					:Any      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )
        UpperCAmelCase_					:Union[str, Any]      			=      BertTokenizerFast.from_pretrained('''bert-base-cased'''     )

        with self.subTest('''BatchEncoding (Python, return_tensors=TENSORFLOW)'''     ):
            self.assert_dump_and_restore(
                tokenizer_p('''Small example to encode'''					,							return_tensors=TensorType.TENSORFLOW     )					,							snake_case     )

        with self.subTest('''BatchEncoding (Rust, return_tensors=TENSORFLOW)'''     ):
            self.assert_dump_and_restore(
                tokenizer_r('''Small example to encode'''					,							return_tensors=TensorType.TENSORFLOW     )					,							snake_case     )





    @require_torch
    @require_tokenizers
    def 				snake_case_ (							self      :  Union[str, Any]     ):
        import torch

        UpperCAmelCase_					:Any      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )
        UpperCAmelCase_					:List[str]      			=      BertTokenizerFast.from_pretrained('''bert-base-cased'''     )

        with self.subTest('''BatchEncoding (Python, return_tensors=PYTORCH)'''     ):
            self.assert_dump_and_restore(
                tokenizer_p('''Small example to encode'''					,							return_tensors=TensorType.PYTORCH     )					,							torch.equal     )

        with self.subTest('''BatchEncoding (Rust, return_tensors=PYTORCH)'''     ):
            self.assert_dump_and_restore(
                tokenizer_r('''Small example to encode'''					,							return_tensors=TensorType.PYTORCH     )					,							torch.equal     )





    @require_tokenizers
    def 				snake_case_ (							self      :  Optional[int]     ):
        UpperCAmelCase_					:str      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )
        UpperCAmelCase_					:Union[str, Any]      			=      BertTokenizerFast.from_pretrained('''bert-base-cased'''     )

        with self.subTest('''Python Tokenizer'''     ):
            self.assertFalse(tokenizer_p('''Small example to_encode'''     ).is_fast     )

        with self.subTest('''Rust Tokenizer'''     ):
            self.assertTrue(tokenizer_r('''Small example to_encode'''     ).is_fast     )





    @require_tokenizers
    def 				snake_case_ (							self      :  Dict     ):
        UpperCAmelCase_					:Tuple      			=      BertTokenizerFast.from_pretrained('''bert-base-cased'''     )
        UpperCAmelCase_					:List[str]      			=      tokenizer_r(['''Test''', '''\xad''', '''test''']					,							is_split_into_words=snake_case     )

        self.assertEqual(encoded.word_to_tokens(0     )					,							TokenSpan(start=1					,							end=2     )     )
        self.assertEqual(encoded.word_to_tokens(1     )					,							snake_case     )
        self.assertEqual(encoded.word_to_tokens(2     )					,							TokenSpan(start=2					,							end=3     )     )





    def 				snake_case_ (							self      :  List[Any]     ):
        UpperCAmelCase_					:int      			=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]}     )
        UpperCAmelCase_					:Union[str, Any]      			=      batch.convert_to_tensors(tensor_type='''np'''     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(2, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(2,)     )
        # test converting the converted
        with CaptureStderr() as cs:
            UpperCAmelCase_					:Optional[Any]      			=      batch.convert_to_tensors(tensor_type='''np'''     )
        self.assertFalse(len(cs.err     )					,							msg=f'should have no warning, but got {cs.err}'     )

        UpperCAmelCase_					:Optional[int]      			=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0}     )
        UpperCAmelCase_					:List[Any]      			=      batch.convert_to_tensors(tensor_type='''np'''					,							prepend_batch_axis=snake_case     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(1, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(1,)     )





    @require_torch
    def 				snake_case_ (							self      :  Union[str, Any]     ):
        UpperCAmelCase_					:List[Any]      			=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]}     )
        UpperCAmelCase_					:List[str]      			=      batch.convert_to_tensors(tensor_type='''pt'''     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(2, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(2,)     )
        # test converting the converted
        with CaptureStderr() as cs:
            UpperCAmelCase_					:str      			=      batch.convert_to_tensors(tensor_type='''pt'''     )
        self.assertFalse(len(cs.err     )					,							msg=f'should have no warning, but got {cs.err}'     )

        UpperCAmelCase_					:Dict      			=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0}     )
        UpperCAmelCase_					:Union[str, Any]      			=      batch.convert_to_tensors(tensor_type='''pt'''					,							prepend_batch_axis=snake_case     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(1, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(1,)     )





    @require_tf
    def 				snake_case_ (							self      :  Tuple     ):
        UpperCAmelCase_					:Optional[Any]      			=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]}     )
        UpperCAmelCase_					:Union[str, Any]      			=      batch.convert_to_tensors(tensor_type='''tf'''     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(2, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(2,)     )
        # test converting the converted
        with CaptureStderr() as cs:
            UpperCAmelCase_					:Optional[int]      			=      batch.convert_to_tensors(tensor_type='''tf'''     )
        self.assertFalse(len(cs.err     )					,							msg=f'should have no warning, but got {cs.err}'     )

        UpperCAmelCase_					:int      			=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0}     )
        UpperCAmelCase_					:Any      			=      batch.convert_to_tensors(tensor_type='''tf'''					,							prepend_batch_axis=snake_case     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(1, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(1,)     )





    @require_flax
    def 				snake_case_ (							self      :  Union[str, Any]     ):
        UpperCAmelCase_					:str      			=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]}     )
        UpperCAmelCase_					:Any      			=      batch.convert_to_tensors(tensor_type='''jax'''     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(2, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(2,)     )
        # test converting the converted
        with CaptureStderr() as cs:
            UpperCAmelCase_					:Union[str, Any]      			=      batch.convert_to_tensors(tensor_type='''jax'''     )
        self.assertFalse(len(cs.err     )					,							msg=f'should have no warning, but got {cs.err}'     )

        UpperCAmelCase_					:str      			=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0}     )
        UpperCAmelCase_					:List[str]      			=      batch.convert_to_tensors(tensor_type='''jax'''					,							prepend_batch_axis=snake_case     )
        self.assertEqual(tensor_batch['''inputs'''].shape					,							(1, 3)     )
        self.assertEqual(tensor_batch['''labels'''].shape					,							(1,)     )





    def 				snake_case_ (							self      :  Optional[int]     ):
        UpperCAmelCase_					:str      			=      [{'''input_ids''': np.array([0, 1, 2]     )}, {'''input_ids''': np.array([0, 1, 2, 3]     )}]
        UpperCAmelCase_					:Any      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )

        UpperCAmelCase_					:List[str]      			=      tokenizer.pad(snake_case					,							padding=snake_case     )
        self.assertTrue(isinstance(batch['''input_ids''']					,							np.ndarray     )     )
        self.assertEqual(batch['''input_ids'''].tolist()					,							[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
        UpperCAmelCase_					:List[str]      			=      tokenizer.pad(snake_case					,							padding=snake_case					,							return_tensors='''np'''     )
        self.assertTrue(isinstance(batch['''input_ids''']					,							np.ndarray     )     )
        self.assertEqual(batch['''input_ids'''].tolist()					,							[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )





    @require_torch
    def 				snake_case_ (							self      :  List[Any]     ):
        import torch

        UpperCAmelCase_					:int      			=      [{'''input_ids''': torch.tensor([0, 1, 2]     )}, {'''input_ids''': torch.tensor([0, 1, 2, 3]     )}]
        UpperCAmelCase_					:List[str]      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )

        UpperCAmelCase_					:int      			=      tokenizer.pad(snake_case					,							padding=snake_case     )
        self.assertTrue(isinstance(batch['''input_ids''']					,							torch.Tensor     )     )
        self.assertEqual(batch['''input_ids'''].tolist()					,							[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
        UpperCAmelCase_					:Union[str, Any]      			=      tokenizer.pad(snake_case					,							padding=snake_case					,							return_tensors='''pt'''     )
        self.assertTrue(isinstance(batch['''input_ids''']					,							torch.Tensor     )     )
        self.assertEqual(batch['''input_ids'''].tolist()					,							[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )





    @require_tf
    def 				snake_case_ (							self      :  Union[str, Any]     ):
        import tensorflow as tf

        UpperCAmelCase_					:List[str]      			=      [{'''input_ids''': tf.constant([0, 1, 2]     )}, {'''input_ids''': tf.constant([0, 1, 2, 3]     )}]
        UpperCAmelCase_					:int      			=      BertTokenizer.from_pretrained('''bert-base-cased'''     )

        UpperCAmelCase_					:str      			=      tokenizer.pad(snake_case					,							padding=snake_case     )
        self.assertTrue(isinstance(batch['''input_ids''']					,							tf.Tensor     )     )
        self.assertEqual(batch['''input_ids'''].numpy().tolist()					,							[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
        UpperCAmelCase_					:List[str]      			=      tokenizer.pad(snake_case					,							padding=snake_case					,							return_tensors='''tf'''     )
        self.assertTrue(isinstance(batch['''input_ids''']					,							tf.Tensor     )     )
        self.assertEqual(batch['''input_ids'''].numpy().tolist()					,							[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )





    @require_tokenizers
    def 				snake_case_ (							self      :  str     ):
        UpperCAmelCase_					:str      			=      Tokenizer(WordPiece(unk_token='''[UNK]'''     )     )
        PreTrainedTokenizerFast(tokenizer_object=snake_case     )





    @require_tokenizers
    def 				snake_case_ (							self      :  str     ):
        UpperCAmelCase_					:Union[str, Any]      			=      Tokenizer(WordPiece(unk_token='''[UNK]'''     )     )
        with tempfile.TemporaryDirectory() as tmpdirname:
            bert_tokenizer.save(os.path.join(snake_case					,							'''tokenizer.json'''     )     )
            PreTrainedTokenizerFast(tokenizer_file=os.path.join(snake_case					,							'''tokenizer.json'''     )     )








"""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


