Remove compiler disable

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -184,9 +184,7 @@ def pytorch_funcify_OpFromGraph(op, node, **kwargs):
     # Apply inner rewrites
     PYTORCH.optimizer(op.fgraph)
     fgraph_fn = pytorch_funcify(op.fgraph, **kwargs, squeeze_output=True)
-    # Disable one step inlining to prevent torch from trying to import local functions
-    # defined in `pytorch_funcify`
-    return torch.compiler.disable(fgraph_fn, recursive=False)
+    return fgraph_fn
 
 
 @pytorch_funcify.register(TensorFromScalar)

pytensor/link/pytorch/dispatch/blockwise.py

@@ -1,5 +1,4 @@
 import torch
-import torch.compiler
 
 from pytensor.graph import FunctionGraph
 from pytensor.link.pytorch.dispatch import pytorch_funcify
@@ -16,7 +15,6 @@ def funcify_Blockwise(op: Blockwise, node, *args, **kwargs):
     for _ in range(batched_dims):
         inner_func = torch.vmap(inner_func)
 
-    @torch.compiler.disable(recursive=False)
     def batcher(*inputs):
         op._check_runtime_broadcast(node, inputs)
         # broadcast on batched_dims

tests/link/pytorch/test_blockwise.py

@@ -29,7 +29,6 @@ def perform(self, *_):
 
 @basic.pytorch_funcify.register(TestOp)
 def evaluate_test_op(op, **_):
-    @torch.compiler.disable(recursive=False)
     def func(a, b):
         op.call_shapes.extend(map(torch.Tensor.size, [a, b]))
         return a @ b