


import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()


def 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase  ):


			'''simple docstring'''







			if gpta_config_file == "":
						__lowercase 					=   GPTaConfig()
			else:
						__lowercase 					=   GPTaConfig.from_json_file(_UpperCamelCase  )
			__lowercase 					=   GPTaModel(_UpperCamelCase  )

			# Load weights from numpy
			load_tf_weights_in_gpta(_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase  )

			# Save pytorch-model
			__lowercase 					=   pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
			__lowercase 					=   pytorch_dump_folder_path + '''/''' + CONFIG_NAME
			print(F'Save PyTorch model to {pytorch_weights_dump_path}'  )
			torch.save(model.state_dict()       ,			_UpperCamelCase  )
			print(F'Save configuration file to {pytorch_config_dump_path}'  )
			with open(_UpperCamelCase       ,			'''w'''       ,			encoding='''utf-8'''  ) as f:
						f.write(config.to_json_string()  )


if __name__ == "__main__":
	a : Optional[Any]		=				argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--gpt2_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
	)
	parser.add_argument(
	    '''--gpt2_config_file''',
	    default='''''',
	    type=str,
	    help=(
	        '''An optional config json file corresponding to the pre-trained OpenAI model. \n'''
	        '''This specifies the model architecture.'''
	    ),
	)
	a : Any		=				parser.parse_args()
	convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)






import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
	import jax
@require_flax
class 	lowerCamelCase_       (				lowerCAmelCase__     ,     unittest.TestCase				):



					'''simple docstring'''







					__UpperCAmelCase            =				FlaxAutoencoderKL






					@property
					def A      (				self						)   ->			Optional[int]:







								'''simple docstring'''

								__lowercase 					=   4
								__lowercase 					=   3
								__lowercase 					=   (3_2, 3_2)

								__lowercase 					=   jax.random.PRNGKey(0						)
								__lowercase 					=   jax.random.uniform(snake_case_ ,       ((batch_size, num_channels) + sizes)						)

								return {"sample": image, "prng_key": prng_key}






					def A      (				self						)   ->			Any:







								'''simple docstring'''

								__lowercase 					=   {
								    '''block_out_channels''': [3_2, 6_4],
								    '''in_channels''': 3,
								    '''out_channels''': 3,
								    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
								    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
								    '''latent_channels''': 4,
								}
								__lowercase 					=   self.dummy_input
								return init_dict, inputs_dict


