def memoize_stacks(func):
 cache = {}
 @wraps(func)
 def memoizer(be, optree):
  (optree_key, tensor_index_map, index_tensor_map) = optree.intrinsic_key_maps()
  optree_key = (optree_key, id(be))
  if (optree_key in cache):
   (stacks, cached_tensor_index_map) = cache[optree_key]
   for stack in stacks:
    for i in range(len(stack)):
     if isinstance(stack[i], Tensor):
      if (stack[i] in cached_tensor_index_map):
       stack[i] = index_tensor_map[cached_tensor_index_map[stack[i]]]
   cache[optree_key] = (stacks, tensor_index_map)
  else:
   stacks = func(be, optree)
   cache[optree_key] = (stacks, tensor_index_map)
  return stacks
 return memoizer

def memoize_stacks ( func ) : cache = { } @ wraps ( func ) def memoizer ( be , optree ) : ( optree_key , tensor_index_map , index_tensor_map ) = optree . intrinsic_key_maps ( ) optree_key = ( optree_key , id ( be ) ) if ( optree_key in cache ) : ( stacks , cached_tensor_index_map ) = cache [ optree_key ] for stack in stacks : for i in range ( len ( stack ) ) : if isinstance ( stack [ i ] , Tensor ) : if ( stack [ i ] in cached_tensor_index_map ) : stack [ i ] = index_tensor_map [ cached_tensor_index_map [ stack [ i ] ] ] cache [ optree_key ] = ( stacks , tensor_index_map ) else : stacks = func ( be , optree ) cache [ optree_key ] = ( stacks , tensor_index_map ) return stacks return memoizer
memoize the stacks using intrinsic_key_maps .