@@ -989,34 +989,45 @@ def apply(self, **kwargs):
989989 # Get items expected to be specialized
990990 specialize = as_tuple (kwargs .pop ('specialize' , []))
991991
992- if not specialize :
993- # Compile the operator before building the arguments list
994- # to avoid out of memory with greedy compilers
995- cfunction = self .cfunction
996-
997- # Build the arguments list to invoke the kernel function
998- with self ._profiler .timer_on ('arguments-preprocess' ):
999- args = self .arguments (** kwargs )
1000- with switch_log_level (comm = args .comm ):
1001- self ._emit_args_profiling ('arguments-preprocess' )
1002-
1003992 # In the case of specialization, arguments must be processed before
1004993 # the operator can be compiled
1005994 if specialize :
1006995 # FIXME: Cannot cope with things like sizes/strides yet since it only
1007996 # looks at the parameters
1008- specialized_args = {p : sympify (args .pop (p .name ))
1009- for p in self .parameters if p .name in specialize }
1010997
1011- op = Specializer (specialized_args ).visit (self )
998+ # Build the arguments list for specialization
999+ with self ._profiler .timer_on ('specialized-arguments-preprocess' ):
1000+ args = self .arguments (** kwargs )
1001+ with switch_log_level (comm = args .comm ):
1002+ self ._emit_args_profiling ('specialized-arguments-preprocess' )
10121003
1013- specialized_kwargs = {k : v for k , v in kwargs .items ()
1014- if k not in specialize }
1004+ # Uses parameters here since Specializer needs {symbol: sympy value} mapper
1005+ specialized_values = {p : sympify (args [p .name ])
1006+ for p in self .parameters if p .name in specialize }
1007+
1008+ op = Specializer (specialized_values ).visit (self )
10151009
10161010 # TODO: Does this cause problems for profilers?
10171011 # FIXME: Need some way to inspect this Operator for testing
10181012 # FIXME: Perhaps this should use some separate method
1019- return op .apply (** specialized_kwargs )
1013+ unspecialized_kwargs = {k : v for k , v in kwargs .items ()
1014+ if k not in specialize }
1015+
1016+ return op .apply (** unspecialized_kwargs )
1017+
1018+ # Compile the operator before building the arguments list
1019+ # to avoid out of memory with greedy compilers
1020+ cfunction = self .cfunction
1021+
1022+ # Build the arguments list to invoke the kernel function
1023+ with self ._profiler .timer_on ('arguments-preprocess' ):
1024+ args = self .arguments (** kwargs )
1025+ with switch_log_level (comm = args .comm ):
1026+ self ._emit_args_profiling ('arguments-preprocess' )
1027+
1028+ args_string = ", " .join ([f"{ p .name } ={ args [p .name ]} "
1029+ for p in self .parameters if p .is_Symbol ])
1030+ debug (f"Invoking `{ self .name } ` with scalar arguments: { args_string } " )
10201031
10211032 # Invoke kernel function with args
10221033 arg_values = [args [p .name ] for p in self .parameters ]
0 commit comments