


'''simple docstring'''




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		):





     def __init__( self       :      List[str]    )				->  int:
        super().__init__()
        lowerCAmelCase     =			nn.Linear(3    ,      4    )
        lowerCAmelCase     =			nn.BatchNormad(4    )
        lowerCAmelCase     =			nn.Linear(4    ,      5    )





     def 							__UpperCAmelCase      ( self       :      List[str]    ,      UpperCAmelCase__       :      Tuple    )				->  Any:
        return self.lineara(self.batchnorm(self.lineara(__lowercase    )    )    )




class 		UpperCAmelCase_      (  UpperCAmelCase__		):





     def 							__UpperCAmelCase      ( self       :      List[str]    ,      UpperCAmelCase__       :      Dict    ,      *UpperCAmelCase__       :      List[str]    ,      **UpperCAmelCase__       :      List[str]    )				->  Tuple:
        return (args[0] + 1,) + args[1:], kwargs




class 		UpperCAmelCase_      (  UpperCAmelCase__		):





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    ,      UpperCAmelCase__       :      Dict    ,      UpperCAmelCase__       :      Union[str, Any]    )				->  str:
        return output + 1




class 		UpperCAmelCase_      (  unittest.TestCase		):





     def 							__UpperCAmelCase      ( self       :      Dict    )				->  Union[str, Any]:
        lowerCAmelCase     =			ModelForTest()
        lowerCAmelCase     =			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 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  int:
        lowerCAmelCase     =			ModelForTest()
        lowerCAmelCase     =			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 							__UpperCAmelCase      ( self       :      Any    )				->  Optional[Any]:
        lowerCAmelCase     =			ModelForTest()
        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			test_model(x + 1    )
        lowerCAmelCase     =			test_model(x + 2    )

        lowerCAmelCase     =			PreForwardHook()
        add_hook_to_module(__lowercase    ,      __lowercase    )
        lowerCAmelCase     =			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
        lowerCAmelCase     =			PreForwardHook()
        add_hook_to_module(__lowercase    ,      __lowercase    )
        lowerCAmelCase     =			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
        lowerCAmelCase     =			SequentialHook(PreForwardHook()    ,      PreForwardHook()    )
        add_hook_to_module(__lowercase    ,      __lowercase    )

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





     def 							__UpperCAmelCase      ( self       :      List[str]    )				->  List[str]:
        lowerCAmelCase     =			ModelForTest()
        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			test_model(__lowercase    )

        lowerCAmelCase     =			PostForwardHook()
        add_hook_to_module(__lowercase    ,      __lowercase    )
        lowerCAmelCase     =			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
        lowerCAmelCase     =			PostForwardHook()
        add_hook_to_module(__lowercase    ,      __lowercase    )
        lowerCAmelCase     =			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
        lowerCAmelCase     =			SequentialHook(PostForwardHook()    ,      PostForwardHook()    )
        add_hook_to_module(__lowercase    ,      __lowercase    )

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





     def 							__UpperCAmelCase      ( self       :      str    )				->  List[str]:
        lowerCAmelCase     =			ModelForTest()
        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			test_model(__lowercase    )

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

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





     @require_multi_gpu
     def 							__UpperCAmelCase      ( self       :      Any    )				->  Any:
        lowerCAmelCase     =			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
        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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    )    )
        lowerCAmelCase     =			torch.randn(2    ,      3    ).to(0    )
        lowerCAmelCase     =			model(__lowercase    )
        self.assertEqual(output.device    ,      torch.device(0    )    )





     def 							__UpperCAmelCase      ( self       :      Dict    )				->  Dict:
        lowerCAmelCase     =			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
        lowerCAmelCase     =			{'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
        lowerCAmelCase     =			torch.device(hook_kwargs['execution_device']    )
        self.assertEqual(model.batchnorm.running_mean.device    ,      __lowercase    )

        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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
        lowerCAmelCase     =			{
            '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'    )    )

        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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 							__UpperCAmelCase      ( self       :      Dict    )				->  Optional[Any]:
        lowerCAmelCase     =			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
        lowerCAmelCase     =			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
        lowerCAmelCase     =			torch.device(__lowercase    )
        self.assertEqual(model.batchnorm.running_mean.device    ,      __lowercase    )

        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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'    )    )

        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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 							__UpperCAmelCase      ( self       :      List[str]    )				->  Any:
        lowerCAmelCase     =			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
        lowerCAmelCase     =			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
        lowerCAmelCase     =			torch.device(__lowercase    )
        self.assertEqual(model.batchnorm.running_mean.device    ,      __lowercase    )

        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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'    )    )

        lowerCAmelCase     =			torch.randn(2    ,      3    )
        lowerCAmelCase     =			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])
