More robust shape check for grad fallback in jacobian

Author: jessegrabowski

pytensor/gradient.py

@@ -2069,13 +2069,13 @@ def jacobian(expression, wrt, consider_constant=None, disconnected_inputs="raise
     else:
         wrt = [wrt]
 
-    if expression.ndim == 0:
+    if all(expression.type.broadcastable):
         # expression is just a scalar, use grad
         return as_list_or_tuple(
             using_list,
             using_tuple,
             grad(
-                expression,
+                expression.squeeze(),
                 wrt,
                 consider_constant=consider_constant,
                 disconnected_inputs=disconnected_inputs,