Skip to content

Commit 5ebbb1a

Browse files
author
Ian Schweer
committed
Allow scipy scalar handling
1 parent d17d4a9 commit 5ebbb1a

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

pytensor/link/pytorch/dispatch/elemwise.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,15 @@
1010
def pytorch_funcify_Elemwise(op, node, **kwargs):
1111
scalar_op = op.scalar_op
1212
base_fn = pytorch_funcify(scalar_op, node=node, **kwargs)
13-
14-
if hasattr(scalar_op, "nfunc_spec") and hasattr(torch, scalar_op.nfunc_spec[0]):
13+
if hasattr(scalar_op, "nfunc_spec") and (
14+
hasattr(torch, scalar_op.nfunc_spec[0]) or "scipy." in scalar_op.nfunc_spec[0]
15+
):
1516
# torch can handle this scalar
1617
# broadcast, we'll let it.
1718
def elemwise_fn(*inputs):
1819
Elemwise._check_runtime_broadcast(node, inputs)
1920
return base_fn(*inputs)
21+
2022
else:
2123

2224
def elemwise_fn(*inputs):

0 commit comments

Comments
 (0)