import atexit
import functools

from cupy import cuda


_memoized_funcs = []


def memoize(for_each_device=False):
    """Makes a function memoizing the result for each argument and device.

    This decorator provides automatic memoization of the function result.

    Args:
        for_each_device (bool): If True, it memoizes the results for each
            device. Otherwise, it memoizes the results only based on the
            arguments.

    """
    def decorator(f):
        global _memoized_funcs
        f._cupy_memo = {}
        _memoized_funcs.append(f)

        @functools.wraps(f)
        def ret(*args, **kwargs):
            arg_key = (args, frozenset(kwargs.items()))
            if for_each_device:
                arg_key = (cuda.Device().id, arg_key)

            memo = f._cupy_memo
            result = memo.get(arg_key, None)
            if result is None:
                result = f(*args, **kwargs)
                memo[arg_key] = result
            return result

        return ret

    return decorator


@atexit.register
def clear_memo():
    """Clears the memoized results for all functions decorated by memoize."""
    global _memoized_funcs
    for f in _memoized_funcs:
        del f._cupy_memo
    _memoized_funcs = []

import atexit
import functools

from cupy import cuda


_memos = []


def memoize(for_each_device=False):
    """Makes a function memoizing the result for each argument and device.

    This decorator provides automatic memoization of the function result.

    Args:
        for_each_device (bool): If True, it memoizes the results for each
            device. Otherwise, it memoizes the results only based on the
            arguments.

    """
    def decorator(f):
        memo = {}
        _memos.append(memo)

        @functools.wraps(f)
        def ret(*args, **kwargs):
            arg_key = (args, frozenset(kwargs.items()))
            if for_each_device:
                arg_key = (cuda.Device().id, arg_key)

            result = memo.get(arg_key, None)
            if result is None:
                result = f(*args, **kwargs)
                memo[arg_key] = result
            return result

        return ret

    return decorator


@atexit.register
def clear_memo():
    """Clears the memoized results for all functions decorated by memoize."""
    for memo in _memos:
        memo.clear()

Fix unintended late finalization of memoized functions
Fix unintended late finalization of memoized functions

  import atexit
  import functools
  
  from cupy import cuda
  
  
- _memoized_funcs = []
+ _memos = []
  
  
  def memoize(for_each_device=False):
      """Makes a function memoizing the result for each argument and device.
  
      This decorator provides automatic memoization of the function result.
  
      Args:
          for_each_device (bool): If True, it memoizes the results for each
              device. Otherwise, it memoizes the results only based on the
              arguments.
  
      """
      def decorator(f):
-         global _memoized_funcs
-         f._cupy_memo = {}
+         memo = {}
-         _memoized_funcs.append(f)
+         _memos.append(memo)
  
          @functools.wraps(f)
          def ret(*args, **kwargs):
              arg_key = (args, frozenset(kwargs.items()))
              if for_each_device:
                  arg_key = (cuda.Device().id, arg_key)
  
-             memo = f._cupy_memo
              result = memo.get(arg_key, None)
              if result is None:
                  result = f(*args, **kwargs)
                  memo[arg_key] = result
              return result
  
          return ret
  
      return decorator
  
  
  @atexit.register
  def clear_memo():
      """Clears the memoized results for all functions decorated by memoize."""
+     for memo in _memos:
+         memo.clear()
-     global _memoized_funcs
-     for f in _memoized_funcs:
-         del f._cupy_memo
-     _memoized_funcs = []
  
Fix unintended late finalization of memoized functions
## Code Before:
import atexit
import functools

from cupy import cuda


_memoized_funcs = []


def memoize(for_each_device=False):
    """Makes a function memoizing the result for each argument and device.

    This decorator provides automatic memoization of the function result.

    Args:
        for_each_device (bool): If True, it memoizes the results for each
            device. Otherwise, it memoizes the results only based on the
            arguments.

    """
    def decorator(f):
        global _memoized_funcs
        f._cupy_memo = {}
        _memoized_funcs.append(f)

        @functools.wraps(f)
        def ret(*args, **kwargs):
            arg_key = (args, frozenset(kwargs.items()))
            if for_each_device:
                arg_key = (cuda.Device().id, arg_key)

            memo = f._cupy_memo
            result = memo.get(arg_key, None)
            if result is None:
                result = f(*args, **kwargs)
                memo[arg_key] = result
            return result

        return ret

    return decorator


@atexit.register
def clear_memo():
    """Clears the memoized results for all functions decorated by memoize."""
    global _memoized_funcs
    for f in _memoized_funcs:
        del f._cupy_memo
    _memoized_funcs = []

## Instruction:
Fix unintended late finalization of memoized functions
## Code After:
import atexit
import functools

from cupy import cuda


_memos = []


def memoize(for_each_device=False):
    """Makes a function memoizing the result for each argument and device.

    This decorator provides automatic memoization of the function result.

    Args:
        for_each_device (bool): If True, it memoizes the results for each
            device. Otherwise, it memoizes the results only based on the
            arguments.

    """
    def decorator(f):
        memo = {}
        _memos.append(memo)

        @functools.wraps(f)
        def ret(*args, **kwargs):
            arg_key = (args, frozenset(kwargs.items()))
            if for_each_device:
                arg_key = (cuda.Device().id, arg_key)

            result = memo.get(arg_key, None)
            if result is None:
                result = f(*args, **kwargs)
                memo[arg_key] = result
            return result

        return ret

    return decorator


@atexit.register
def clear_memo():
    """Clears the memoized results for all functions decorated by memoize."""
    for memo in _memos:
        memo.clear()
