




import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu
class 		UpperCAmelCase   (				nn.Module			):







						'''simple docstring'''





						def __init__(						self							:     List[str]						):






							"""simple docstring"""






							super().__init__()
							snake_case_		    =     nn.Linear(3       ,	4						)
							snake_case_		    =     nn.BatchNormad(4						)
							snake_case_		    =     nn.Linear(4       ,	5						)






						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     Tuple						):






							"""simple docstring"""






							return self.lineara(self.batchnorm(self.lineara(__lowercase						)						)						)







class 		UpperCAmelCase   (				UpperCAmelCase__			):







						'''simple docstring'''





						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     Dict       ,	*__lowercase							:     List[str]       ,	**__lowercase							:     List[str]						):






							"""simple docstring"""






							return (args[0] + 1,) + args[1:], kwargs







class 		UpperCAmelCase   (				UpperCAmelCase__			):







						'''simple docstring'''





						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     Dict       ,	__lowercase							:     Union[str, Any]						):






							"""simple docstring"""






							return output + 1







class 		UpperCAmelCase   (				unittest.TestCase			):







						'''simple docstring'''





						def        snake_case__						(						self							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()
							snake_case_		    =     ModelHook()

							add_hook_to_module(__lowercase       ,	__lowercase						)
							self.assertEqual(test_model._hf_hook       ,	__lowercase						)
							self.assertTrue(hasattr(__lowercase       ,	"_old_forward"						)						)

							# Check adding the hook did not change the name or the signature
							self.assertEqual(test_model.forward.__name__       ,	"forward"						)
							self.assertListEqual(list(inspect.signature(test_model.forward						).parameters						)       ,	["x"]						)

							remove_hook_from_module(__lowercase						)
							self.assertFalse(hasattr(__lowercase       ,	"_hf_hook"						)						)
							self.assertFalse(hasattr(__lowercase       ,	"_old_forward"						)						)






						def        snake_case__						(						self							:     Optional[Any]						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()
							snake_case_		    =     ModelHook()

							add_hook_to_module(__lowercase       ,	__lowercase						)
							add_hook_to_module(__lowercase       ,	__lowercase       ,	append=__lowercase						)

							self.assertEqual(isinstance(test_model._hf_hook       ,	__lowercase						)       ,	__lowercase						)
							self.assertEqual(len(test_model._hf_hook.hooks						)       ,	2						)
							self.assertTrue(hasattr(__lowercase       ,	"_old_forward"						)						)

							# Check adding the hook did not change the name or the signature
							self.assertEqual(test_model.forward.__name__       ,	"forward"						)
							self.assertListEqual(list(inspect.signature(test_model.forward						).parameters						)       ,	["x"]						)

							remove_hook_from_module(__lowercase						)
							self.assertFalse(hasattr(__lowercase       ,	"_hf_hook"						)						)
							self.assertFalse(hasattr(__lowercase       ,	"_old_forward"						)						)






						def        snake_case__						(						self							:     Any						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()
							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     test_model(x + 1						)
							snake_case_		    =     test_model(x + 2						)

							snake_case_		    =     PreForwardHook()
							add_hook_to_module(__lowercase       ,	__lowercase						)
							snake_case_		    =     test_model(__lowercase						)
							self.assertTrue(torch.allclose(__lowercase       ,	__lowercase       ,	atol=1E-5						)						)

							# Attaching a hook to a model when it already has one replaces, does not chain
							snake_case_		    =     PreForwardHook()
							add_hook_to_module(__lowercase       ,	__lowercase						)
							snake_case_		    =     test_model(__lowercase						)
							self.assertTrue(torch.allclose(__lowercase       ,	__lowercase       ,	atol=1E-5						)						)

							# You need to use the sequential hook to chain two or more hooks
							snake_case_		    =     SequentialHook(PreForwardHook()       ,	PreForwardHook()						)
							add_hook_to_module(__lowercase       ,	__lowercase						)

							snake_case_		    =     test_model(__lowercase						)
							assert torch.allclose(__lowercase       ,	__lowercase       ,	atol=1E-5						)






						def        snake_case__						(						self							:     List[str]						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()
							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     test_model(__lowercase						)

							snake_case_		    =     PostForwardHook()
							add_hook_to_module(__lowercase       ,	__lowercase						)
							snake_case_		    =     test_model(__lowercase						)
							self.assertTrue(torch.allclose(__lowercase       ,	output + 1       ,	atol=1E-5						)						)

							# Attaching a hook to a model when it already has one replaces, does not chain
							snake_case_		    =     PostForwardHook()
							add_hook_to_module(__lowercase       ,	__lowercase						)
							snake_case_		    =     test_model(__lowercase						)
							self.assertTrue(torch.allclose(__lowercase       ,	output + 1       ,	atol=1E-5						)						)

							# You need to use the sequential hook to chain two or more hooks
							snake_case_		    =     SequentialHook(PostForwardHook()       ,	PostForwardHook()						)
							add_hook_to_module(__lowercase       ,	__lowercase						)

							snake_case_		    =     test_model(__lowercase						)
							assert torch.allclose(__lowercase       ,	output + 2       ,	atol=1E-5						)






						def        snake_case__						(						self							:     str						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()
							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     test_model(__lowercase						)

							snake_case_		    =     PostForwardHook()
							add_hook_to_module(__lowercase       ,	__lowercase						)
							snake_case_		    =     test_model(__lowercase						)
							self.assertTrue(torch.allclose(__lowercase       ,	output + 1						)						)
							self.assertTrue(outputa.requires_grad						)

							snake_case_		    =     True
							snake_case_		    =     test_model(__lowercase						)
							self.assertFalse(outputa.requires_grad						)






						@require_multi_gpu
						def        snake_case__						(						self							:     Any						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()
							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# This will move each submodule on different devices
							add_hook_to_module(model.lineara       ,	AlignDevicesHook(execution_device=0						)						)
							add_hook_to_module(model.batchnorm       ,	AlignDevicesHook(execution_device=0						)						)
							add_hook_to_module(model.lineara       ,	AlignDevicesHook(execution_device=1						)						)

							self.assertEqual(model.lineara.weight.device       ,	torch.device(0						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device(0						)						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	torch.device(0						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device(1						)						)

							# We can still make a forward pass. The input does not need to be on any particular device
							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	torch.device(1						)						)

							# We can add a general hook to put back output on same device as input.
							add_hook_to_module(__lowercase       ,	AlignDevicesHook(io_same_device=__lowercase						)						)
							snake_case_		    =     torch.randn(2       ,	3						).to(0						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	torch.device(0						)						)






						def        snake_case__						(						self							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# This will move each submodule on different devices
							snake_case_		    =     {"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

							add_hook_to_module(model.lineara       ,	AlignDevicesHook(**__lowercase						)						)
							add_hook_to_module(model.batchnorm       ,	AlignDevicesHook(**__lowercase						)						)
							add_hook_to_module(model.lineara       ,	AlignDevicesHook(**__lowercase						)						)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							# Buffers are not included in the offload by default, so are on the execution device
							snake_case_		    =     torch.device(hook_kwargs["execution_device"]						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	__lowercase						)

							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	__lowercase						)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_module(model.lineara						)
							remove_hook_from_module(model.batchnorm						)
							remove_hook_from_module(model.lineara						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# Now test with buffers included in the offload
							snake_case_		    =     {
							    "execution_device": 0 if torch.cuda.is_available() else "cpu",
							    "offload": True,
							    "offload_buffers": True,
							}

							add_hook_to_module(model.lineara       ,	AlignDevicesHook(**__lowercase						)						)
							add_hook_to_module(model.batchnorm       ,	AlignDevicesHook(**__lowercase						)						)
							add_hook_to_module(model.lineara       ,	AlignDevicesHook(**__lowercase						)						)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	torch.device("meta"						)						)

							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	__lowercase						)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_module(model.lineara						)
							remove_hook_from_module(model.batchnorm						)
							remove_hook_from_module(model.lineara						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)






						def        snake_case__						(						self							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# This will move each submodule on different devices
							snake_case_		    =     0 if torch.cuda.is_available() else "cpu"
							attach_align_device_hook(__lowercase       ,	execution_device=__lowercase       ,	offload=__lowercase						)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							# Buffers are not included in the offload by default, so are on the execution device
							snake_case_		    =     torch.device(__lowercase						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	__lowercase						)

							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	__lowercase						)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(__lowercase						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# Now test with buffers included in the offload
							attach_align_device_hook(__lowercase       ,	execution_device=__lowercase       ,	offload=__lowercase       ,	offload_buffers=__lowercase						)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	torch.device("meta"						)						)

							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	__lowercase						)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(__lowercase						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)






						def        snake_case__						(						self							:     List[str]						):






							"""simple docstring"""






							snake_case_		    =     ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# This will move each submodule on different devices
							snake_case_		    =     0 if torch.cuda.is_available() else "cpu"
							attach_align_device_hook(
							    __lowercase       ,	execution_device=__lowercase       ,	offload=__lowercase       ,	weights_map=model.state_dict()						)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							# Buffers are not included in the offload by default, so are on the execution device
							snake_case_		    =     torch.device(__lowercase						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	__lowercase						)

							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	__lowercase						)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(__lowercase						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)

							# Now test with buffers included in the offload
							attach_align_device_hook(
							    __lowercase       ,	execution_device=__lowercase       ,	offload=__lowercase       ,	weights_map=model.state_dict()       ,	offload_buffers=__lowercase       ,	)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("meta"						)						)
							self.assertEqual(model.batchnorm.running_mean.device       ,	torch.device("meta"						)						)

							snake_case_		    =     torch.randn(2       ,	3						)
							snake_case_		    =     model(__lowercase						)
							self.assertEqual(output.device       ,	__lowercase						)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(__lowercase						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.batchnorm.weight.device       ,	torch.device("cpu"						)						)
							self.assertEqual(model.lineara.weight.device       ,	torch.device("cpu"						)						)






from __future__ import annotations

lowercase__				:     str 		=						[
    [-1, 0],  # left
    [0, -1],  # down
    [1, 0],  # right
    [0, 1],  # up
]


def 						lowerCamelCase__    (		_A     ,	_A     ,	_A     ,	_A     ,	_A     ,	):







	'''simple docstring'''





	snake_case_		    =     [
	    [0 for col in range(len(grid[0]							)							)] for row in range(len(_A							)							)
	]  # the reference grid
	snake_case_		    =     1
	snake_case_		    =     [
	    [0 for col in range(len(grid[0]							)							)] for row in range(len(_A							)							)
	]  # the action grid

	snake_case_		    =     init[0]
	snake_case_		    =     init[1]
	snake_case_		    =     0
	snake_case_		    =     g + heuristic[x][y]  # cost from starting cell to destination cell
	snake_case_		    =     [[f, g, x, y]]

	snake_case_		    =     False  # flag that is set when search is complete
	snake_case_		    =     False  # flag set if we can't find expand

	while not found and not resign:
		if len(_A							) == 0:
			raise ValueError("Algorithm is unable to find solution"							)
		else:  # to choose the least costliest action so as to move closer to the goal
			cell.sort()
			cell.reverse()
			snake_case_		    =     cell.pop()
			snake_case_		    =     next_cell[2]
			snake_case_		    =     next_cell[3]
			snake_case_		    =     next_cell[1]

			if x == goal[0] and y == goal[1]:
				snake_case_		    =     True
			else:
				for i in range(len(_A							)							):  # to try out different valid actions
					snake_case_		    =     x + DIRECTIONS[i][0]
					snake_case_		    =     y + DIRECTIONS[i][1]
					if xa >= 0 and xa < len(_A							) and ya >= 0 and ya < len(grid[0]							):
						if closed[xa][ya] == 0 and grid[xa][ya] == 0:
							snake_case_		    =     g + cost
							snake_case_		    =     ga + heuristic[xa][ya]
							cell.append([fa, ga, xa, ya]							)
							snake_case_		    =     1
							snake_case_		    =     i
	snake_case_		    =     []
	snake_case_		    =     goal[0]
	snake_case_		    =     goal[1]
	invpath.append([x, y]							)  # we get the reverse path from here
	while x != init[0] or y != init[1]:
		snake_case_		    =     x - DIRECTIONS[action[x][y]][0]
		snake_case_		    =     y - DIRECTIONS[action[x][y]][1]
		snake_case_		    =     xa
		snake_case_		    =     ya
		invpath.append([x, y]							)

	snake_case_		    =     []
	for i in range(len(_A							)							):
		path.append(invpath[len(_A							) - 1 - i]							)
	return path, action


if __name__ == "__main__":
		lowercase__				:     Union[str, Any] 		=						[
		    [0, 1, 0, 0, 0, 0],
		    [0, 1, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
		    [0, 1, 0, 0, 0, 0],
		    [0, 1, 0, 0, 1, 0],
		    [0, 0, 0, 0, 1, 0],
		]

		lowercase__				:     Optional[Any] 		=						[0, 0]
		# all coordinates are given in format [y,x]
		lowercase__				:     Tuple 		=						[len(grid) - 1, len(grid[0]) - 1]
		lowercase__				:     Dict 		=						1

		# the cost map which pushes the path closer to the goal
		lowercase__				:     Any 		=						[[0 for row in range(len(grid[0]))] for col in range(len(grid))]
		for i in range(len(grid)):
				for j in range(len(grid[0])):
						lowercase__				:     Union[str, Any] 		=						abs(i - goal[0]) + abs(j - goal[1])
						if grid[i][j] == 1:
								# added extra penalty in the heuristic map
								lowercase__				:     int 		=						99

		lowercase__		,       lowercase__				:     Tuple 		=						search(grid, init, goal, cost, heuristic)

		print("ACTION MAP")
		for i in range(len(action)):
				print(action[i])

		for i in range(len(path)):
				print(path[i])
