Skip to content

Commit 0f0f57b

Browse files
committed
Simplify Blockwise perform method
1 parent 70fb494 commit 0f0f57b

File tree

1 file changed

+6
-5
lines changed

1 file changed

+6
-5
lines changed

pytensor/tensor/blockwise.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -317,7 +317,7 @@ def core_func(
317317
):
318318
self.core_op.perform(
319319
core_node,
320-
[np.asarray(inp) for inp in inner_inputs],
320+
inner_inputs,
321321
inner_outputs_storage,
322322
)
323323

@@ -326,7 +326,11 @@ def core_func(
326326
else:
327327
return tuple(r[0] for r in inner_outputs_storage)
328328

329-
gufunc = np.vectorize(core_func, signature=self.signature)
329+
gufunc = np.vectorize(
330+
core_func,
331+
signature=self.signature,
332+
otypes=[out.type.dtype for out in node.outputs],
333+
)
330334

331335
node.tag.gufunc = gufunc
332336

@@ -361,9 +365,6 @@ def perform(self, node, inputs, output_storage):
361365
res = (res,)
362366

363367
for node_out, out_storage, r in zip(node.outputs, output_storage, res):
364-
out_dtype = getattr(node_out, "dtype", None)
365-
if out_dtype and out_dtype != r.dtype:
366-
r = np.asarray(r, dtype=out_dtype)
367368
out_storage[0] = r
368369

369370
def __str__(self):

0 commit comments

Comments
 (0)