



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


 '''simple docstring'''




 def __init__(   self) ->      str:



      """simple docstring"""







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




 def       lowercase_  (   self							,  lowerCAmelCase_) ->      Union[str, Any]:



      """simple docstring"""







      return self.lineara(self.batchnorm(self.lineara(lowerCAmelCase_)))







class 	UpperCAmelCase	(		__a):


 '''simple docstring'''




 def       lowercase_  (   self							,  lowerCAmelCase_							,  *lowerCAmelCase_							,  **lowerCAmelCase_) ->      Optional[Any]:



      """simple docstring"""







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







class 	UpperCAmelCase	(		__a):


 '''simple docstring'''




 def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      Optional[Any]:



      """simple docstring"""







      return output + 1







class 	UpperCAmelCase	(		unittest.TestCase):


 '''simple docstring'''




 def       lowercase_  (   self) ->      Union[str, Any]:



      """simple docstring"""







      a_       		=ModelForTest()
      a_       		=ModelHook()

      add_hook_to_module(lowerCAmelCase_							,  lowerCAmelCase_)
      self.assertEqual(test_model._hf_hook							,  lowerCAmelCase_)
      self.assertTrue(hasattr(lowerCAmelCase_							,  "_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(lowerCAmelCase_)
      self.assertFalse(hasattr(lowerCAmelCase_							,  "_hf_hook"))
      self.assertFalse(hasattr(lowerCAmelCase_							,  "_old_forward"))




 def       lowercase_  (   self) ->      str:



      """simple docstring"""







      a_       		=ModelForTest()
      a_       		=ModelHook()

      add_hook_to_module(lowerCAmelCase_							,  lowerCAmelCase_)
      add_hook_to_module(lowerCAmelCase_							,  lowerCAmelCase_							,  append=lowerCAmelCase_)

      self.assertEqual(isinstance(test_model._hf_hook							,  lowerCAmelCase_)							,  lowerCAmelCase_)
      self.assertEqual(len(test_model._hf_hook.hooks)							,  2)
      self.assertTrue(hasattr(lowerCAmelCase_							,  "_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(lowerCAmelCase_)
      self.assertFalse(hasattr(lowerCAmelCase_							,  "_hf_hook"))
      self.assertFalse(hasattr(lowerCAmelCase_							,  "_old_forward"))




 def       lowercase_  (   self) ->      Optional[int]:



      """simple docstring"""







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

      a_       		=PreForwardHook()
      add_hook_to_module(lowerCAmelCase_							,  lowerCAmelCase_)
      a_       		=test_model(lowerCAmelCase_)
      self.assertTrue(torch.allclose(lowerCAmelCase_							,  lowerCAmelCase_							,  atol=1e-5))

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

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

      a_       		=test_model(lowerCAmelCase_)
      assert torch.allclose(lowerCAmelCase_							,  lowerCAmelCase_							,  atol=1e-5)




 def       lowercase_  (   self) ->      int:



      """simple docstring"""







      a_       		=ModelForTest()
      a_       		=torch.randn(2							,  3)
      a_       		=test_model(lowerCAmelCase_)

      a_       		=PostForwardHook()
      add_hook_to_module(lowerCAmelCase_							,  lowerCAmelCase_)
      a_       		=test_model(lowerCAmelCase_)
      self.assertTrue(torch.allclose(lowerCAmelCase_							,  output + 1							,  atol=1e-5))

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

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

      a_       		=test_model(lowerCAmelCase_)
      assert torch.allclose(lowerCAmelCase_							,  output + 2							,  atol=1e-5)




 def       lowercase_  (   self) ->      Dict:



      """simple docstring"""







      a_       		=ModelForTest()
      a_       		=torch.randn(2							,  3)
      a_       		=test_model(lowerCAmelCase_)

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

      a_       		=True
      a_       		=test_model(lowerCAmelCase_)
      self.assertFalse(outputa.requires_grad)




 @require_multi_gpu
 def       lowercase_  (   self) ->      str:



      """simple docstring"""







      a_       		=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
      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      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(lowerCAmelCase_							,  AlignDevicesHook(io_same_device=lowerCAmelCase_))
      a_       		=torch.randn(2							,  3).to(0)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  torch.device(0))




 def       lowercase_  (   self) ->      Tuple:



      """simple docstring"""







      a_       		=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
      a_       		={"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

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

      # 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
      a_       		=torch.device(hook_kwargs["execution_device"])
      self.assertEqual(model.batchnorm.running_mean.device							,  lowerCAmelCase_)

      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  lowerCAmelCase_)

      # 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
      a_       		={
          "execution_device": 0 if torch.cuda.is_available() else "cpu",
          "offload": True,
          "offload_buffers": True,
      }

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

      # 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"))

      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  lowerCAmelCase_)

      # 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       lowercase_  (   self) ->      str:



      """simple docstring"""







      a_       		=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
      a_       		=0 if torch.cuda.is_available() else "cpu"
      attach_align_device_hook(lowerCAmelCase_							,  execution_device=lowerCAmelCase_							,  offload=lowerCAmelCase_)

      # 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
      a_       		=torch.device(lowerCAmelCase_)
      self.assertEqual(model.batchnorm.running_mean.device							,  lowerCAmelCase_)

      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  lowerCAmelCase_)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCAmelCase_)
      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(lowerCAmelCase_							,  execution_device=lowerCAmelCase_							,  offload=lowerCAmelCase_							,  offload_buffers=lowerCAmelCase_)

      # 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"))

      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  lowerCAmelCase_)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCAmelCase_)
      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       lowercase_  (   self) ->      List[str]:



      """simple docstring"""







      a_       		=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
      a_       		=0 if torch.cuda.is_available() else "cpu"
      attach_align_device_hook(
          lowerCAmelCase_							,  execution_device=lowerCAmelCase_							,  offload=lowerCAmelCase_							,  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
      a_       		=torch.device(lowerCAmelCase_)
      self.assertEqual(model.batchnorm.running_mean.device							,  lowerCAmelCase_)

      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  lowerCAmelCase_)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCAmelCase_)
      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(
          lowerCAmelCase_							,  execution_device=lowerCAmelCase_							,  offload=lowerCAmelCase_							,  weights_map=model.state_dict()							,  offload_buffers=lowerCAmelCase_							,  )

      # 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"))

      a_       		=torch.randn(2							,  3)
      a_       		=model(lowerCAmelCase_)
      self.assertEqual(output.device							,  lowerCAmelCase_)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCAmelCase_)
      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"))







'''simple docstring'''


import os
from math import logaa






def  UpperCAmelCase_       (		lowercase__ = "base_exp.txt"       ):


     '''simple docstring'''







     a_       		=0
     a_       		=0
     for i, line in enumerate(open(os.path.join(os.path.dirname(lowercase__       )       ,	lowercase__       )       )       ):
          a_		,						a_       		=list(map(lowercase__       ,	line.split(","       )       )       )
          if x * logaa(lowercase__       ) > largest:
               a_       		=x * logaa(lowercase__       )
               a_       		=i + 1
     return result


if __name__ == "__main__":
  print(solution())


