






"""simple docstring"""
from __future__ import annotations






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   ):
	A_		:						Any	    =   len(snake_case__					)

	# If row is equal to the size of the board it means there are a queen in each row in
	# the current board (possible_board)
	if row == n:
		# We convert the variable possible_board that looks like this: [1, 3, 0, 2] to
		# this: ['. Q . . ', '. . . Q ', 'Q . . . ', '. . Q . ']
		boards.append([""". """ * i + """Q """ + """. """ * (n - 1 - i) for i in possible_board]					)
		return

	# We iterate each column in the row to find all possible results in each row
	for col in range(snake_case__					):
		# We apply that we learned previously. First we check that in the current board
		# (possible_board) there are not other same value because if there is it means
		# that there are a collision in vertical. Then we apply the two formulas we
		# learned before:
		#
		# 45º: y - x = b or 45: row - col = b
		# 135º: y + x = b or row + col = b.
		#
		# And we verify if the results of this two formulas not exist in their variables
		# respectively.  (diagonal_right_collisions, diagonal_left_collisions)
		#
		# If any or these are True it means there is a collision so we continue to the
		# next value in the for loop.
		if (
		    col in possible_board
		    or row - col in diagonal_right_collisions
		    or row + col in diagonal_left_collisions
		):
			continue

		# If it is False we call dfs function again and we update the inputs
		depth_first_search(
		    [*possible_board, col]   ,   [*diagonal_right_collisions, row - col]   ,   [*diagonal_left_collisions, row + col]   ,   snake_case__   ,   snake_case__   ,   )






def 			__UpperCamelCase	(		snake_case__					):
	A_		:						list[list[str]]	    =   []
	depth_first_search([]   ,   []   ,   []   ,   snake_case__   ,   snake_case__					)

	# Print all the boards
	for board in boards:
		for column in board:
			print(snake_case__					)
		print(""""""					)

	print(len(snake_case__					)   ,   """solutions were found."""					)


if __name__ == "__main__":
							import doctest

							doctest.testmod()
							n_queens_solution(4)









"""simple docstring"""
import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version(">=", FSDP_PYTORCH_VERSION):
							import torch.distributed.checkpoint as dist_cp
							from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
							from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
							from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
							from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_lowerCAmelCase         =							get_logger(__name__)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		A_		:						Optional[int]	    =   model.state_dict()
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A_		:						int	    =   F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
			A_		:						Union[str, Any]	    =   os.path.join(snake_case__   ,   snake_case__					)
			if accelerator.process_index == 0:
				logger.info(F"""Saving model to {output_model_file}"""					)
				torch.save(snake_case__   ,   snake_case__					)
				logger.info(F"""Model saved to {output_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A_		:						List[str]	    =   (
			    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
			    if model_index == 0
			    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
			)
			A_		:						List[str]	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Saving model to {output_model_file}"""					)
			torch.save(snake_case__   ,   snake_case__					)
			logger.info(F"""Model saved to {output_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A_		:						str	    =   os.path.join(snake_case__   ,   F"""{MODEL_NAME}_{model_index}"""					)
			os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
			logger.info(F"""Saving model to {ckpt_dir}"""					)
			A_		:						Dict	    =   {"""model""": state_dict}

			dist_cp.save_state_dict(
			    state_dict=snake_case__   ,   storage_writer=dist_cp.FileSystemWriter(snake_case__					)   ,   planner=DefaultSavePlanner()   ,   )
			logger.info(F"""Model saved to {ckpt_dir}"""					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if type(snake_case__					) != FSDP and accelerator.process_index != 0:
				if not fsdp_plugin.sync_module_states:
					raise ValueError(
					    """Set the `sync_module_states` flag to `True` so that model states are synced across processes when """
					    """initializing FSDP object"""					)
				return
			A_		:						Tuple	    =   F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
			A_		:						Any	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Loading model from {input_model_file}"""					)
			A_		:						List[Any]	    =   torch.load(snake_case__					)
			logger.info(F"""Model loaded from {input_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A_		:						List[str]	    =   (
			    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
			    if model_index == 0
			    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
			)
			A_		:						Any	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Loading model from {input_model_file}"""					)
			A_		:						List[str]	    =   torch.load(snake_case__					)
			logger.info(F"""Model loaded from {input_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A_		:						Optional[Any]	    =   (
			    os.path.join(snake_case__   ,   F"""{MODEL_NAME}_{model_index}"""					)
			    if F"""{MODEL_NAME}""" not in input_dir
			    else input_dir
			)
			logger.info(F"""Loading model from {ckpt_dir}"""					)
			A_		:						List[str]	    =   {"""model""": model.state_dict()}
			dist_cp.load_state_dict(
			    state_dict=snake_case__   ,   storage_reader=dist_cp.FileSystemReader(snake_case__					)   ,   planner=DefaultLoadPlanner()   ,   )
			A_		:						Tuple	    =   state_dict["""model"""]
			logger.info(F"""Model loaded from {ckpt_dir}"""					)
		model.load_state_dict(snake_case__					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		A_		:						Optional[Any]	    =   FSDP.optim_state_dict(snake_case__   ,   snake_case__					)
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if accelerator.process_index == 0:
				A_		:						Any	    =   (
				    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
				)
				A_		:						Optional[int]	    =   os.path.join(snake_case__   ,   snake_case__					)
				logger.info(F"""Saving Optimizer state to {output_optimizer_file}"""					)
				torch.save(snake_case__   ,   snake_case__					)
				logger.info(F"""Optimizer state saved in {output_optimizer_file}"""					)
		else:
			A_		:						Union[str, Any]	    =   os.path.join(snake_case__   ,   F"""{OPTIMIZER_NAME}_{optimizer_index}"""					)
			os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
			logger.info(F"""Saving Optimizer state to {ckpt_dir}"""					)
			dist_cp.save_state_dict(
			    state_dict={"""optimizer""": optim_state}   ,   storage_writer=dist_cp.FileSystemWriter(snake_case__					)   ,   planner=DefaultSavePlanner()   ,   )
			logger.info(F"""Optimizer state saved in {ckpt_dir}"""					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A_		:						List[str]	    =   None
			# below check should work but currently it isn't working (mostly opytorch issue),
			# in the meantime disabling it at the cost of excess memory usage
			# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
			A_		:						Tuple	    =   (
			    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
			)
			A_		:						int	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Loading Optimizer state from {input_optimizer_file}"""					)
			A_		:						Tuple	    =   torch.load(snake_case__					)
			logger.info(F"""Optimizer state loaded from {input_optimizer_file}"""					)
		else:
			A_		:						str	    =   (
			    os.path.join(snake_case__   ,   F"""{OPTIMIZER_NAME}_{optimizer_index}"""					)
			    if F"""{OPTIMIZER_NAME}""" not in input_dir
			    else input_dir
			)
			logger.info(F"""Loading Optimizer from {ckpt_dir}"""					)
			A_		:						Union[str, Any]	    =   load_sharded_optimizer_state_dict(
			    model_state_dict=model.state_dict()   ,   optimizer_key="""optimizer"""   ,   storage_reader=dist_cp.FileSystemReader(snake_case__					)   ,   )
			A_		:						Tuple	    =   optim_state["""optimizer"""]
			logger.info(F"""Optimizer loaded from {ckpt_dir}"""					)
		A_		:						List[Any]	    =   FSDP.optim_state_dict_to_load(snake_case__   ,   snake_case__   ,   snake_case__					)
		optimizer.load_state_dict(snake_case__					)

