code: 
  
  
  
  
  
  import importlib.metadata
  import subprocess
  import sys
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Tuple = False				) ->    Optional[Any]:
    __snake_case   :				List[Any]        =       False
    if "IPython" in sys.modules:
      __snake_case   :				Dict        =       "google.colab" in str(sys.modules["IPython"].get_ipython()				)
  
    if in_colab:
      if upgrade:
        __snake_case   :				Dict        =       ["pip", "install", "-U", "torch"]
        subprocess.run(UpperCamelCase__   ,     check=UpperCamelCase__				)
      # get the current version of torch
      __snake_case   :				Any        =       importlib.metadata.version("torch"				)
      __snake_case   :				Optional[Any]        =       torch_version[: torch_version.rindex("."				)]
      __snake_case   :				Union[str, Any]        =       f"""https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl"""
      __snake_case   :				Any        =       ["pip", "install", xla_wheel]
      subprocess.run(UpperCamelCase__   ,     check=UpperCamelCase__				)
    else:
      raise RuntimeError("`install_xla` utility works only on google colab."				)
  
  
style_context: 
  
  
  
  
  
  import torch
  from torch import autograd
  
  
  
  
  class 				_lowerCamelCase (    autograd.Function    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        @staticmethod
        def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   str:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Dict        =       inputs.numel()
          __snake_case   :				Optional[int]        =       int(0.005 * nb_elems   ) + 1
          if sigmoid:
            __snake_case   :				Any        =       (torch.sigmoid(UpperCAmelCase   ) > threshold).type(inputs.type()   )
          else:
            __snake_case   :				Optional[int]        =       (inputs > threshold).type(inputs.type()   )
          if mask.sum() < nb_min:
            # We limit the pruning so that at least 0.5% (half a percent) of the weights are remaining
            __snake_case   :				List[str]        =       inputs.flatten().kthvalue(max(nb_elems - nb_min				,						1   )   ).values
            __snake_case   :				Union[str, Any]        =       (inputs > k_threshold).type(inputs.type()   )
          return mask
  
  
  
  
  
  
  
        @staticmethod
        def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase   )       ->   Optional[int]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          return gradOutput, None, None
  
  
  
  
  class 				_lowerCamelCase (    autograd.Function    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        @staticmethod
        def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   Optional[int]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Optional[int]        =       inputs.clone()
          __snake_case    ,   __snake_case   :				int        =       inputs.flatten().sort(descending=UpperCAmelCase   )
          __snake_case   :				Tuple        =       int(threshold * inputs.numel()   )
  
          # flat_out and mask access the same memory.
          __snake_case   :				Optional[int]        =       mask.flatten()
          __snake_case   :				Union[str, Any]        =       0
          __snake_case   :				Any        =       1
          return mask
  
  
  
  
  
  
  
        @staticmethod
        def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase   )       ->   List[str]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          return gradOutput, None
  
  
  
  
  class 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        @staticmethod
        def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Union[str, Any]        =       inputs.clone()
          __snake_case    ,   __snake_case   :				Optional[Any]        =       inputs.abs().flatten().sort(descending=UpperCAmelCase   )
          __snake_case   :				List[Any]        =       int(threshold * inputs.numel()   )
  
          # flat_out and mask access the same memory.
          __snake_case   :				Optional[Any]        =       mask.flatten()
          __snake_case   :				Tuple        =       0
          __snake_case   :				Any        =       1
          return mask
  
  