






from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	torch.Tensor:
    # The first operation in a checkpoint can't be in-place, but it's
    # nice to have in-place addition during inference. Thus...
    if not inplace:
        SCREAMING_SNAKE_CASE__       							= ma + ma
    else:
        ma += ma

    return ma


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	torch.Tensor:
    SCREAMING_SNAKE_CASE__       							= -1 * len(lowerCAmelCase_  )
    SCREAMING_SNAKE_CASE__       							= list(range(len(tensor.shape[:zero_index]  )  )  )
    return tensor.permute(first_inds + [zero_index + i for i in inds]  )


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	torch.Tensor:
    return t.reshape(t.shape[:-no_dims] + (-1,)  )


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_ = 1e-4  )							->	torch.Tensor:
    SCREAMING_SNAKE_CASE__       							= mask.expand(*value.shape  )
    return torch.sum(mask * value   ,       dim=lowerCAmelCase_  ) / (eps + torch.sum(lowerCAmelCase_   ,       dim=lowerCAmelCase_  ))


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_ = 2.31_25   ,       lowerCAmelCase_ = 21.68_75   ,       lowerCAmelCase_ = 6_4  )							->	torch.Tensor:
    SCREAMING_SNAKE_CASE__       							= torch.linspace(lowerCAmelCase_   ,       lowerCAmelCase_   ,       no_bins - 1   ,       device=pts.device  )
    SCREAMING_SNAKE_CASE__       							= torch.sqrt(torch.sum((pts.unsqueeze(-2  ) - pts.unsqueeze(-3  )) ** 2   ,       dim=-1  )  )
    return torch.bucketize(lowerCAmelCase_   ,       lowerCAmelCase_  )


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	dict:
    SCREAMING_SNAKE_CASE__       							= dicts[0]
    SCREAMING_SNAKE_CASE__       							= {}
    for k, v in first.items():
        SCREAMING_SNAKE_CASE__       							= [d[k] for d in dicts]
        if isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ):
            SCREAMING_SNAKE_CASE__       							= dict_multimap(lowerCAmelCase_   ,       lowerCAmelCase_  )
        else:
            SCREAMING_SNAKE_CASE__       							= fn(lowerCAmelCase_  )

    return new_dict


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	torch.Tensor:
    SCREAMING_SNAKE_CASE__       							= v_bins.view(((1,) * len(x.shape  )) + (len(lowerCAmelCase_  ),)  )
    SCREAMING_SNAKE_CASE__       							= x[..., None] - reshaped_bins
    SCREAMING_SNAKE_CASE__       							= torch.argmin(torch.abs(lowerCAmelCase_  )   ,       dim=-1  )
    return nn.functional.one_hot(lowerCAmelCase_   ,       num_classes=len(lowerCAmelCase_  )  ).float()


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_ = 0   ,       lowerCAmelCase_ = 0  )							->	torch.Tensor:
    SCREAMING_SNAKE_CASE__       							= []
    for i, s in enumerate(data.shape[:no_batch_dims]  ):
        SCREAMING_SNAKE_CASE__       							= torch.arange(lowerCAmelCase_  )
        SCREAMING_SNAKE_CASE__       							= r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape  ) - i - 1)))  )
        ranges.append(lowerCAmelCase_  )

    SCREAMING_SNAKE_CASE__       							= [slice(lowerCAmelCase_  ) for _ in range(len(data.shape  ) - no_batch_dims  )]
    SCREAMING_SNAKE_CASE__       							= inds
    ranges.extend(lowerCAmelCase_  )
    # Matt note: Editing this to get around the behaviour of using a list as an array index changing
    # in recent Numpy versions
    return data[tuple(lowerCAmelCase_  )]


_A				:							List[str]							= TypeVar("""T""")


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	Dict[Any, Union[dict, list, tuple, Any]]:
    SCREAMING_SNAKE_CASE__       							= {}
    for k, v in dic.items():
        if isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ):
            SCREAMING_SNAKE_CASE__       							= dict_map(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        else:
            SCREAMING_SNAKE_CASE__       							= tree_map(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )

    return new_dict


@overload
def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	Any:
    ...


@overload
def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	dict:
    ...


@overload
def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	list:
    ...


@overload
def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	tuple:
    ...


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	List[str]:
    if isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ):
        return dict_map(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
    elif isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ):
        return [tree_map(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  ) for x in tree]
    elif isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ):
        return tuple(tree_map(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  ) for x in tree  )
    elif isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ):
        return fn(lowerCAmelCase_  )
    else:
        print(type(lowerCAmelCase_  )  )
        raise ValueError('''Not supported'''  )



_A				:							Optional[int]							= partial(tree_map, leaf_type=torch.Tensor)







import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
  import transformer_engine.pytorch as te


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True  )							->	Tuple:
    if not is_fpa_available():
        raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''  )
    for name, module in model.named_children():
        if isinstance(lowerCAmelCase_   ,       nn.Linear  ) and to_transformer_engine and _convert_linear:
            # Return early if the linear layer weights are not multiples of 16
            if any(p % 1_6 != 0 for p in module.weight.shape  ):
                return
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= te.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       nn.LayerNorm  ) and to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= te.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.Linear  ) and not to_transformer_engine and _convert_linear:
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= nn.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.LayerNorm  ) and not to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= nn.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        else:
            convert_model(
                lowerCAmelCase_   ,       to_transformer_engine=lowerCAmelCase_   ,       _convert_linear=_convert_linear   ,       _convert_ln=_convert_ln   ,       )



def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    if not is_fpa_available():
        raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''  )
    for m in model.modules():
        if isinstance(lowerCAmelCase_   ,       (te.LayerNorm, te.Linear)  ):
            return True
    return False
