Remove restrictive check in local_blockwise_alloc

ricardoV94 · ricardoV94 · commit 146a0a84f07f · 2024-05-07T18:37:56.000+02:00
diff --git a/pytensor/tensor/rewriting/blockwise.py b/pytensor/tensor/rewriting/blockwise.py
@@ -188,6 +188,5 @@ def local_blockwise_alloc(fgraph, node):
             )
             for new_out in new_outs
         ]
-    assert new_outs[0].type.broadcastable == old_out_type.broadcastable
     copy_stack_trace(node.outputs, new_outs)
     return new_outs

Original file line number	Diff line number	Diff line change
`@@ -188,6 +188,5 @@ def local_blockwise_alloc(fgraph, node):`
`188`	`188`	`)`
`189`	`189`	`for new_out in new_outs`
`190`	`190`	`]`
`191`		`- assert new_outs[0].type.broadcastable == old_out_type.broadcastable`
`192`	`191`	`copy_stack_trace(node.outputs, new_outs)`
`193`	`192`	`return new_outs`