


'''simple docstring'''
import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase = False     ):






       """simple docstring"""





       if device_specific:
              seed += AcceleratorState().process_index
       random.seed(UpperCAmelCase     )
       np.random.seed(UpperCAmelCase     )
       torch.manual_seed(UpperCAmelCase     )
       if is_xpu_available():
              torch.xpu.manual_seed_all(UpperCAmelCase     )
       elif is_npu_available():
              torch.npu.manual_seed_all(UpperCAmelCase     )
       else:
              torch.cuda.manual_seed_all(UpperCAmelCase     )
       # ^^ safe to call this function even if cuda is not available
       if is_tpu_available():
              xm.set_rng_state(UpperCAmelCase     )



def 				lowerCAmelCase__	(       UpperCAmelCase = None    ,       UpperCAmelCase = None     ):






       """simple docstring"""





       if rng_type == RNGType.TORCH:
              snake_case__					:	Optional[Any]          =					torch.get_rng_state()
       elif rng_type == RNGType.CUDA:
              snake_case__					:	Any          =					torch.cuda.get_rng_state()
       elif rng_type == RNGType.XLA:
              assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
              snake_case__					:	List[str]          =					torch.tensor(xm.get_rng_state()     )
       elif rng_type == RNGType.NPU:
              assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
              snake_case__					:	Dict          =					torch.npu.get_rng_state()
       elif rng_type == RNGType.XPU:
              assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
              snake_case__					:	Optional[Any]          =					torch.xpu.get_rng_state()
       elif rng_type == RNGType.GENERATOR:
              assert generator is not None, "Need a generator to synchronize its seed."
              snake_case__					:	List[Any]          =					generator.get_state()

       # Broadcast the rng state from device 0 to other devices
       snake_case__					:	str          =					AcceleratorState()
       if state.distributed_type == DistributedType.TPU:
              snake_case__					:	Tuple          =					rng_state.to(xm.xla_device()     )
              xm.collective_broadcast([rng_state]     )
              xm.mark_step()
              snake_case__					:	List[str]          =					rng_state.cpu()
       elif (
           state.distributed_type in CUDA_DISTRIBUTED_TYPES
           or state.distributed_type == DistributedType.MULTI_NPU
           or state.distributed_type == DistributedType.MULTI_XPU
       ):
              snake_case__					:	Any          =					rng_state.to(state.device     )
              torch.distributed.broadcast(UpperCAmelCase    ,       0     )
              snake_case__					:	Optional[Any]          =					rng_state.cpu()
       elif state.distributed_type == DistributedType.MULTI_CPU:
              torch.distributed.broadcast(UpperCAmelCase    ,       0     )

       # Set the broadcast rng state
       if rng_type == RNGType.TORCH:
              torch.set_rng_state(UpperCAmelCase     )
       elif rng_type == RNGType.CUDA:
              torch.cuda.set_rng_state(UpperCAmelCase     )
       elif rng_type == RNGType.NPU:
              torch.npu.set_rng_state(UpperCAmelCase     )
       elif rng_type == RNGType.XPU:
              torch.xpu.set_rng_state(UpperCAmelCase     )
       elif rng_type == RNGType.XLA:
              xm.set_rng_state(rng_state.item()     )
       elif rng_type == RNGType.GENERATOR:
              generator.set_state(UpperCAmelCase     )






def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase = None     ):






       """simple docstring"""





       for rng_type in rng_types:
              synchronize_rng_state(RNGType(UpperCAmelCase     )    ,       generator=UpperCAmelCase     )








'''simple docstring'''
from typing import Dict, Optional, Tuple

import torch



def 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





       snake_case__					:	Optional[int]          =					boundaries[1] - boundaries[0]
       snake_case__					:	Optional[Any]          =					boundaries + step / 2
       snake_case__					:	str          =					torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1     )]    ,       dim=0     )
       return bin_centers



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase    ,       ):






       """simple docstring"""





       snake_case__					:	str          =					_calculate_bin_centers(UpperCAmelCase     )
       return (
           torch.sum(aligned_distance_error_probs * bin_centers    ,       dim=-1     ),
           bin_centers[-1],
       )



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase = 31    ,       UpperCAmelCase = 64    ,       **UpperCAmelCase    ,       ):






       """simple docstring"""





       snake_case__					:	Dict          =					torch.linspace(0    ,       UpperCAmelCase    ,       steps=(no_bins - 1)    ,       device=logits.device     )

       snake_case__					:	str          =					torch.nn.functional.softmax(UpperCAmelCase    ,       dim=-1     )
       snake_case__       ,   snake_case__					:	List[str]          =					_calculate_expected_aligned_error(
           alignment_confidence_breaks=UpperCAmelCase    ,       aligned_distance_error_probs=UpperCAmelCase    ,       )

       return {
           "aligned_confidence_probs": aligned_confidence_probs,
           "predicted_aligned_error": predicted_aligned_error,
           "max_predicted_aligned_error": max_predicted_aligned_error,
       }






def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase = None    ,       UpperCAmelCase = 31    ,       UpperCAmelCase = 64    ,       UpperCAmelCase = 1E-8    ,       **UpperCAmelCase    ,       ):






       """simple docstring"""





       if residue_weights is None:
              snake_case__					:	List[Any]          =					logits.new_ones(logits.shape[-2]     )

       snake_case__					:	Optional[Any]          =					torch.linspace(0    ,       UpperCAmelCase    ,       steps=(no_bins - 1)    ,       device=logits.device     )

       snake_case__					:	int          =					_calculate_bin_centers(UpperCAmelCase     )
       torch.sum(UpperCAmelCase     )
       snake_case__					:	Dict          =					logits.shape[-2]
       snake_case__					:	Optional[Any]          =					max(UpperCAmelCase    ,       19     )

       snake_case__					:	Optional[int]          =					1.2_4 * (clipped_n - 15) ** (1.0 / 3) - 1.8

       snake_case__					:	Tuple          =					torch.nn.functional.softmax(UpperCAmelCase    ,       dim=-1     )

       snake_case__					:	int          =					1.0 / (1 + (bin_centers**2) / (da**2))
       snake_case__					:	Optional[int]          =					torch.sum(probs * tm_per_bin    ,       dim=-1     )

       snake_case__					:	Optional[Any]          =					residue_weights / (eps + residue_weights.sum())
       snake_case__					:	Any          =					torch.sum(predicted_tm_term * normed_residue_mask    ,       dim=-1     )

       snake_case__					:	Dict          =					per_alignment * residue_weights

       snake_case__					:	List[str]          =					(weighted == torch.max(UpperCAmelCase     )).nonzero()[0]
       return per_alignment[tuple(UpperCAmelCase     )]




