Defer the use of nfunc_spec in JAX scalar dispatch

Author: ferrine

pytensor/link/jax/dispatch/scalar.py

@@ -62,16 +62,14 @@ def check_if_inputs_scalars(node):
 
 @jax_funcify.register(ScalarOp)
 def jax_funcify_ScalarOp(op, node, **kwargs):
-    func_name = op.nfunc_spec[0]
-
     # We dispatch some PyTensor operators to Python operators
     # whenever the inputs are all scalars.
     are_inputs_scalars = check_if_inputs_scalars(node)
     if are_inputs_scalars:
         elemwise = elemwise_scalar(op)
         if elemwise is not None:
             return elemwise
-
+    func_name = op.nfunc_spec[0]
     if "." in func_name:
         jnp_func = functools.reduce(getattr, [jax] + func_name.split("."))
     else: