





"""simple docstring"""






import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()

def    __lowerCamelCase     (							SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE				)		->							Union[str, Any]:






       """simple docstring"""




       _UpperCAmelCase			     =  FNetConfig.from_json_file(SCREAMING_SNAKE_CASE				)
       print(F"""Building PyTorch model from configuration: {config}"""				)
       _UpperCAmelCase			     =  FNetForPreTraining(SCREAMING_SNAKE_CASE				)

       _UpperCAmelCase			     =  restore_checkpoint(SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE				)
       _UpperCAmelCase			     =  checkpoint_dict['target']

       # Embeddings
       # Position IDs
       _UpperCAmelCase			     =  fnet_pretraining_model.state_dict()

       _UpperCAmelCase			     =  state_dict['fnet.embeddings.position_ids']
       _UpperCAmelCase			     =  {'fnet.embeddings.position_ids': position_ids}
       # Embedding Layers
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['word']['embedding']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['position']['embedding'][0]				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['type']['embedding']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['hidden_mapping_in']['kernel']				).T
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['hidden_mapping_in']['bias']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['layer_norm']['scale']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['layer_norm']['bias']				)

       # Encoder Layers
       for layer in range(config.num_hidden_layers				):
              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['mixing_layer_norm']['scale']				)
              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['mixing_layer_norm']['bias']				)

              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['intermediate']['kernel']				).T
              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['intermediate']['bias']				)

              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['output']['kernel']				).T
              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['output']['bias']				)

              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['output_layer_norm']['scale']				)
              _UpperCAmelCase			     =  torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['output_layer_norm']['bias']				)

       # Pooler Layers
       _UpperCAmelCase			     =  torch.tensor(pretrained_model_params['encoder']['pooler']['kernel']				).T
       _UpperCAmelCase			     =  torch.tensor(pretrained_model_params['encoder']['pooler']['bias']				)

       # Masked LM Layers
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['predictions_dense']['kernel']				).T
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['predictions_dense']['bias']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['predictions_layer_norm']['scale']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['predictions_layer_norm']['bias']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['encoder']['embedder']['word']['embedding']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['predictions_output']['output_bias']				)
       _UpperCAmelCase			     =  torch.tensor(pretrained_model_params['predictions_output']['output_bias']				)

       # Seq Relationship Layers
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['classification']['output_kernel']				)
       _UpperCAmelCase			     =  torch.tensor(
           pretrained_model_params['classification']['output_bias']				)

       # Load State Dict
       fnet_pretraining_model.load_state_dict(SCREAMING_SNAKE_CASE				)

       # Save PreTrained
       print(F"""Saving pretrained model to {save_path}"""				)
       fnet_pretraining_model.save_pretrained(SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
    lowerCAmelCase_      =				argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--flax_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
    )
    parser.add_argument(
        '''--fnet_config_file''',
        default=None,
        type=str,
        required=True,
        help=(
            '''The config json file corresponding to the pre-trained FNet model. \n'''
            '''This specifies the model architecture.'''
        ),
    )
    parser.add_argument('''--save_path''', default=None, type=str, required=True, help='''Path to the output model.''')
    lowerCAmelCase_      =				parser.parse_args()
    convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)













"""simple docstring"""






from collections import defaultdict

def    __lowerCamelCase     (							SCREAMING_SNAKE_CASE = 5				)		->							int:






       """simple docstring"""




       _UpperCAmelCase			     =  defaultdict(SCREAMING_SNAKE_CASE				)
       _UpperCAmelCase			     =  0

       while True:
              _UpperCAmelCase			     =  get_digits(SCREAMING_SNAKE_CASE				)
              freqs[digits].append(SCREAMING_SNAKE_CASE				)

              if len(freqs[digits]				) == max_base:
                     _UpperCAmelCase			     =  freqs[digits][0] ** 3
                     return base

              num += 1

def    __lowerCamelCase     (							SCREAMING_SNAKE_CASE				)		->							str:






       """simple docstring"""




       return "".join(sorted(str(num**3				)				)				)


if __name__ == "__main__":
    print(F'''{solution() = }''')






