diff --git a/pyopencl/tools.py b/pyopencl/tools.py index fc7f459d40ba07547b53a1826feca6d20f61c0ae..d7e56dc9fc05c7365855ecb0b99daaf998a4289c 100644 --- a/pyopencl/tools.py +++ b/pyopencl/tools.py @@ -83,6 +83,10 @@ def first_arg_dependent_memoize(func): .. versionadded:: 2011.2 """ + if kwargs: + cache_key = (args, kwargs) + else: + cache_key = args try: ctx_dict = func._pyopencl_first_arg_dep_memoize_dic except AttributeError: @@ -92,11 +96,11 @@ def first_arg_dependent_memoize(func): _first_arg_dependent_caches.append(ctx_dict) try: - return ctx_dict[cl_object][args] + return ctx_dict[cl_object][cache_key] except KeyError: arg_dict = ctx_dict.setdefault(cl_object, {}) result = func(cl_object, *args, **kwargs) - arg_dict[args] = result + arg_dict[cache_key] = result return result from functools import update_wrapper