



'''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 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      None				):


			'''simple docstring'''







			...



@overload
def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor				):


			'''simple docstring'''







			...



def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      Dict							,							lowerCamelCase__   :      List[str]							,							lowerCamelCase__   :      Dict				):


			'''simple docstring'''







			A:      str														=   aatype == rc.restype_order["""G"""]
			A:      Dict														=   rc.atom_order["""CA"""]
			A:      Dict														=   rc.atom_order["""CB"""]
			A:      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:
						A:      Optional[int]														=   torch.where(
						    lowerCamelCase__							,							all_atom_masks[..., ca_idx]							,							all_atom_masks[..., cb_idx]							,							)
						return pseudo_beta, pseudo_beta_mask
			else:
						return pseudo_beta



def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      Dict[str, torch.Tensor]				):


			'''simple docstring'''







			A:      Union[str, Any]														=   batched_gather(
			    lowerCamelCase__							,							batch["""residx_atom37_to_atom14"""]							,							dim=-2							,							no_batch_dims=len(atomaa.shape[:-2]				)							,							)

			A:      Optional[int]														=   atomaa_data * batch["""atom37_atom_exists"""][..., None]

			return atomaa_data



def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      Dict[str, torch.Tensor]				):


			'''simple docstring'''







			A:      Dict														=   template_feats["""template_aatype"""]
			A:      List[str]														=   template_feats["""template_torsion_angles_sin_cos"""]
			A:      List[str]														=   template_feats["""template_alt_torsion_angles_sin_cos"""]
			A:      List[str]														=   template_feats["""template_torsion_angles_mask"""]
			A:      Optional[int]														=   torch.cat(
			    [
			        nn.functional.one_hot(lowerCamelCase__							,							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 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      Dict[str, torch.Tensor]							,							lowerCamelCase__   :      torch.types.Number							,							lowerCamelCase__   :      torch.types.Number							,							lowerCamelCase__   :      int							,							lowerCamelCase__   :      bool = False							,							lowerCamelCase__   :      float = 1e-2_0							,							lowerCamelCase__   :      float = 1e8							,							):


			'''simple docstring'''







			A:      List[Any]														=   batch["""template_pseudo_beta_mask"""]
			A:      List[str]														=   template_mask[..., None] * template_mask[..., None, :]

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

			A:      Dict														=   [dgram, template_mask_ad[..., None]]

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

			A:      Union[str, Any]														=   batch["""template_aatype"""].shape[-1]
			to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2]							,							lowerCamelCase__							,							-1							,							-1				)				)
			to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2]							,							-1							,							lowerCamelCase__							,							-1				)				)

			A							,		A							,		A:      List[Any]														=   [rc.atom_order[a] for a in ["""N""", """CA""", """C"""]]
			A:      Dict														=   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=lowerCamelCase__							,							)
			A:      Dict														=   rigids.get_trans()[..., None, :, :]
			A:      Optional[int]														=   rigids[..., None].invert_apply(lowerCamelCase__				)

			A:      int														=   torch.rsqrt(eps + torch.sum(rigid_vec**2							,							dim=-1				)				)

			A:      str														=   batch["""template_all_atom_mask"""]
			A:      List[str]														=   t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
			A:      Dict														=   template_mask[..., None] * template_mask[..., None, :]

			A:      str														=   inv_distance_scalar * template_mask_ad
			A:      Union[str, Any]														=   rigid_vec * inv_distance_scalar[..., None]

			if not use_unit_vector:
						A:      List[str]														=   unit_vector * 0.0

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

			A:      Tuple														=   torch.cat(lowerCamelCase__							,							dim=-1				)
			A:      int														=   act * template_mask_ad[..., None]

			return act



def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      Dict[str, torch.Tensor]				):


			'''simple docstring'''







			A:      torch.LongTensor														=   nn.functional.one_hot(batch["""extra_msa"""]							,							23				)
			A:      Optional[int]														=   [
			    msa_ahot,
			    batch["""extra_has_deletion"""].unsqueeze(-1				),
			    batch["""extra_deletion_value"""].unsqueeze(-1				),
			]
			return torch.cat(lowerCamelCase__							,							dim=-1				)



def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      Rigid							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							):


			'''simple docstring'''







			A:      int														=   rrgdf[aatype, ...]

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

			A:      Union[str, Any]														=   alpha.new_zeros((*((1,) * len(alpha.shape[:-1]				)), 2)				)
			A:      Any														=   1

			# [*, N, 8, 2]
			A:      str														=   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.

			A:      Union[str, Any]														=   alpha.new_zeros(default_r.get_rots().get_rot_mats().shape				)
			A:      int														=   1
			A:      Dict														=   alpha[..., 1]
			A:      int														=   -alpha[..., 0]
			A:      str														=   alpha

			A:      str														=   default_r.compose(Rigid(Rotation(rot_mats=lowerCamelCase__				)							,							lowerCamelCase__				)				)

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

			A:      str														=   all_frames[..., 4]
			A:      Dict														=   chia_frame_to_bb.compose(lowerCamelCase__				)
			A:      int														=   chia_frame_to_bb.compose(lowerCamelCase__				)
			A:      int														=   chia_frame_to_bb.compose(lowerCamelCase__				)

			A:      Union[str, Any]														=   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							,							)

			A:      List[Any]														=   r[..., None].compose(lowerCamelCase__				)

			return all_frames_to_global



def 		_SCREAMING_SNAKE_CASE							(				lowerCamelCase__   :      Rigid							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							lowerCamelCase__   :      torch.Tensor							,							):


			'''simple docstring'''







			A:      List[Any]														=   group_idx[aatype, ...]

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

			# [*, N, 14, 8]
			A:      Union[str, Any]														=   r[..., None, :] * group_mask_one_hot

			# [*, N, 14]
			A:      Union[str, Any]														=   t_atoms_to_global.map_tensor_fn(lambda lowerCamelCase__				: torch.sum(lowerCamelCase__							,							dim=-1				)				)

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

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

			return pred_positions










'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__SCREAMING_SNAKE_CASE :      Dict             ={
    'configuration_nystromformer': ['NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'NystromformerConfig'],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__SCREAMING_SNAKE_CASE :      Dict             =[
					    'NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
					    'NystromformerForMaskedLM',
					    'NystromformerForMultipleChoice',
					    'NystromformerForQuestionAnswering',
					    'NystromformerForSequenceClassification',
					    'NystromformerForTokenClassification',
					    'NystromformerLayer',
					    'NystromformerModel',
					    'NystromformerPreTrainedModel',
					]


if TYPE_CHECKING:
					from .configuration_nystromformer import NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, NystromformerConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_nystromformer import (
										    NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    NystromformerForMaskedLM,
										    NystromformerForMultipleChoice,
										    NystromformerForQuestionAnswering,
										    NystromformerForSequenceClassification,
										    NystromformerForTokenClassification,
										    NystromformerLayer,
										    NystromformerModel,
										    NystromformerPreTrainedModel,
										)


else:
					import sys

					__SCREAMING_SNAKE_CASE :      Dict             =_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





