



'''simple docstring'''



import tensorflow as tf

from ...tf_utils import shape_list





class _UpperCamelCase		(     tf.keras.layers.Layer		):





    '''simple docstring'''





    def __init__(							self  :  List[Any]  ,       _lowerCAmelCase  :  Any  ,       _lowerCAmelCase  :  Any  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  str  ,       _lowerCAmelCase  :  Tuple=1  ,       _lowerCAmelCase  :  Optional[int]=False  ,       **_lowerCAmelCase  :  Optional[Any]):







           '''simple docstring'''

           super().__init__(**A_)

           __lowercase       	=vocab_size
           __lowercase       	=d_embed
           __lowercase       	=d_proj

           __lowercase       	=cutoffs + [vocab_size]
           __lowercase       	=[0] + self.cutoffs
           __lowercase       	=div_val

           __lowercase       	=self.cutoffs[0]
           __lowercase       	=len(self.cutoffs) - 1
           __lowercase       	=self.shortlist_size + self.n_clusters
           __lowercase       	=keep_order

           __lowercase       	=[]
           __lowercase       	=[]

    def 				__lowerCamelCase     (							self  :  Optional[Any]  ,       _lowerCAmelCase  :  Dict):







           '''simple docstring'''

           if self.n_clusters > 0:
                  __lowercase       	=self.add_weight(
                      shape=(self.n_clusters, self.d_embed)  ,       initializer='zeros'  ,       trainable=A_  ,       name='cluster_weight')
                  __lowercase       	=self.add_weight(
                      shape=(self.n_clusters,)  ,       initializer='zeros'  ,       trainable=A_  ,       name='cluster_bias')

           if self.div_val == 1:
                  for i in range(len(self.cutoffs)):
                         if self.d_proj != self.d_embed:
                                __lowercase       	=self.add_weight(
                                    shape=(self.d_embed, self.d_proj)  ,       initializer='zeros'  ,       trainable=A_  ,       name=f"""out_projs_._{i}"""  ,       )
                                self.out_projs.append(A_)
                         else:
                                self.out_projs.append(A_)
                         __lowercase       	=self.add_weight(
                             shape=(self.vocab_size, self.d_embed)  ,       initializer='zeros'  ,       trainable=A_  ,       name=f"""out_layers_._{i}_._weight"""  ,       )
                         __lowercase       	=self.add_weight(
                             shape=(self.vocab_size,)  ,       initializer='zeros'  ,       trainable=A_  ,       name=f"""out_layers_._{i}_._bias"""  ,       )
                         self.out_layers.append((weight, bias))
           else:
                  for i in range(len(self.cutoffs)):
                         __lowercase	,				__lowercase       	=self.cutoff_ends[i], self.cutoff_ends[i + 1]
                         __lowercase       	=self.d_embed // (self.div_val**i)

                         __lowercase       	=self.add_weight(
                             shape=(d_emb_i, self.d_proj)  ,       initializer='zeros'  ,       trainable=A_  ,       name=f"""out_projs_._{i}""")
                         self.out_projs.append(A_)
                         __lowercase       	=self.add_weight(
                             shape=(r_idx - l_idx, d_emb_i)  ,       initializer='zeros'  ,       trainable=A_  ,       name=f"""out_layers_._{i}_._weight"""  ,       )
                         __lowercase       	=self.add_weight(
                             shape=(r_idx - l_idx,)  ,       initializer='zeros'  ,       trainable=A_  ,       name=f"""out_layers_._{i}_._bias"""  ,       )
                         self.out_layers.append((weight, bias))
           super().build(A_)

    @staticmethod
    def 				__lowerCamelCase     (							_lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  List[str]  ,       _lowerCAmelCase  :  Optional[int]=None):







           '''simple docstring'''

           __lowercase       	=x
           if proj is not None:
                  __lowercase       	=tf.einsum('ibd,ed->ibe'  ,       A_  ,       A_)
           return tf.einsum('ibd,nd->ibn'  ,       A_  ,       A_) + b

    @staticmethod
    def 				__lowerCamelCase     (							_lowerCAmelCase  :  Optional[Any]  ,       _lowerCAmelCase  :  Optional[int]):







           '''simple docstring'''

           __lowercase       	=shape_list(A_)
           __lowercase       	=tf.range(lp_size[0]  ,       dtype=target.dtype)
           __lowercase       	=tf.stack([r, target]  ,       1)
           return tf.gather_nd(A_  ,       A_)

    def 				__lowerCamelCase     (							self  :  int  ,       _lowerCAmelCase  :  List[str]  ,       _lowerCAmelCase  :  List[Any]  ,       _lowerCAmelCase  :  str=True  ,       _lowerCAmelCase  :  Any=False):







           '''simple docstring'''

           __lowercase       	=0
           if self.n_clusters == 0:
                  __lowercase       	=self._logit(A_  ,       self.out_layers[0][0]  ,       self.out_layers[0][1]  ,       self.out_projs[0])
                  if target is not None:
                         __lowercase       	=tf.nn.sparse_softmax_cross_entropy_with_logits(labels=A_  ,       logits=A_)
                  __lowercase       	=tf.nn.log_softmax(A_  ,       axis=-1)
           else:
                  __lowercase       	=shape_list(A_)
                  __lowercase       	=[]
                  __lowercase       	=tf.zeros(hidden_sizes[:2])
                  for i in range(len(self.cutoffs)):
                         __lowercase	,				__lowercase       	=self.cutoff_ends[i], self.cutoff_ends[i + 1]
                         if target is not None:
                                __lowercase       	=(target >= l_idx) & (target < r_idx)
                                __lowercase       	=tf.where(A_)
                                __lowercase       	=tf.boolean_mask(A_  ,       A_) - l_idx

                         if self.div_val == 1:
                                __lowercase       	=self.out_layers[0][0][l_idx:r_idx]
                                __lowercase       	=self.out_layers[0][1][l_idx:r_idx]
                         else:
                                __lowercase       	=self.out_layers[i][0]
                                __lowercase       	=self.out_layers[i][1]

                         if i == 0:
                                __lowercase       	=tf.concat([cur_W, self.cluster_weight]  ,       0)
                                __lowercase       	=tf.concat([cur_b, self.cluster_bias]  ,       0)

                                __lowercase       	=self._logit(A_  ,       A_  ,       A_  ,       self.out_projs[0])
                                __lowercase       	=tf.nn.log_softmax(A_)
                                out.append(head_logprob[..., : self.cutoffs[0]])
                                if target is not None:
                                       __lowercase       	=tf.boolean_mask(A_  ,       A_)
                                       __lowercase       	=self._gather_logprob(A_  ,       A_)
                         else:
                                __lowercase       	=self._logit(A_  ,       A_  ,       A_  ,       self.out_projs[i])
                                __lowercase       	=tf.nn.log_softmax(A_)
                                __lowercase       	=self.cutoffs[0] + i - 1  # No probability for the head cluster
                                __lowercase       	=head_logprob[..., cluster_prob_idx, None] + tail_logprob
                                out.append(A_)
                                if target is not None:
                                       __lowercase       	=tf.boolean_mask(A_  ,       A_)
                                       __lowercase       	=tf.boolean_mask(A_  ,       A_)
                                       __lowercase       	=self._gather_logprob(A_  ,       A_)
                                       cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                         if target is not None:
                                loss += tf.scatter_nd(A_  ,       -cur_logprob  ,       shape_list(A_))
                  __lowercase       	=tf.concat(A_  ,       axis=-1)

           if target is not None:
                  if return_mean:
                         __lowercase       	=tf.reduce_mean(A_)
                  # Add the training-time loss value to the layer using `self.add_loss()`.
                  self.add_loss(A_)

                  # Log the loss as a metric (we could log arbitrary metrics,
                  # including different metrics for training and inference.
                  self.add_metric(A_  ,       name=self.name  ,       aggregation='mean' if return_mean else '')

           return out








'''simple docstring'''



# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os

from accelerate.utils import ComputeEnvironment

from .cluster import get_cluster_input
from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
from .sagemaker import get_sagemaker_input


lowerCamelCase	 =       """Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine"""



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=_ask_options(
							    'In which compute environment are you running?'  ,   ['This machine', 'AWS (Amazon SageMaker)']  ,   _convert_compute_environment  ,   )
							if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
														__lowercase       	=get_sagemaker_input()
							else:
														__lowercase       	=get_cluster_input()
							return config



def   _A       (   _lowerCAmelCase=None					):



							"""simple docstring"""



							if subparsers is not None:
														__lowercase       	=subparsers.add_parser('config'  ,   description=_lowerCAmelCase					)
							else:
														__lowercase       	=argparse.ArgumentParser('Accelerate config command'  ,   description=_lowerCAmelCase					)

							parser.add_argument(
							    '--config_file'  ,   default=_lowerCAmelCase  ,   help=(
							        'The path to use to store the config file. Will default to a file named default_config.yaml in the cache '
							        'location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '
							        'such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '
							        'with \'huggingface\'.'
							    )  ,   )

							if subparsers is not None:
														parser.set_defaults(func=_lowerCAmelCase					)
							return parser



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=get_user_input()
							if args.config_file is not None:
														__lowercase       	=args.config_file
							else:
														if not os.path.isdir(_lowerCAmelCase					):
																					os.makedirs(_lowerCAmelCase					)
														__lowercase       	=default_yaml_config_file

							if config_file.endswith('.json'					):
														config.to_json_file(_lowerCAmelCase					)
							else:
														config.to_yaml_file(_lowerCAmelCase					)
							print(f"""accelerate configuration saved at {config_file}"""					)



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=config_command_parser()
							__lowercase       	=parser.parse_args()
							config_command(_lowerCAmelCase					)


if __name__ == "__main__":
		main()



