fix(convexadam): disable only the TensorExpr fuser during registration

The TS feature models trip torch's TensorExpr fuser alias analysis (INTERNAL ASSERT on aten::size/dim in FuseTensorExprs) when itk-impact runs them. Verified via itk-impact: model-side is impossible (itk-impact passes an undefined direction, so guarding it needs a shape op = the fuser trigger); disabling the fuser is the only fix. Scope it (_no_texpr_fuser) to the registration: modern profiling executor stays on (not legacy), measured ~1% cost (fuser only touches the few feature forwards, not the 150-iter C++ optimisation loop).

ConvexAdam_Coarse/Model.py

@@ -33,6 +33,7 @@ NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engi
 runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
 """
 
+import contextlib
 import json
 import os
 from dataclasses import dataclass, field
@@ -106,6 +107,22 @@ class ModelSpec:
     distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
 
 
+@contextlib.contextmanager
+def _no_texpr_fuser():
+    """Disable ONLY torch's TensorExpr (NNC) fuser for the block. The TS feature-model graphs have a direction/
+    orientation branch whose shape ops (aten::dim / aten::size) crash the fuser's alias analysis under itk-impact
+    ("INTERNAL ASSERT ... We don't have an op for aten::size" in FuseTensorExprs). It cannot be fixed model-side:
+    itk-impact passes an *undefined* direction, so guarding it needs a shape op — the very fuser trigger. The
+    modern profiling executor stays on (NOT the legacy executor); measured cost ~1% of a registration (the fuser
+    only touches the few feature forwards, not the C++ optimisation loop). The caller's setting is restored."""
+    prev = torch._C._jit_texpr_fuser_enabled()
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    try:
+        yield
+    finally:
+        torch._C._jit_set_texpr_fuser_enabled(prev)
+
+
 def _coarse_registration_type():
     """The coupled-convex initializer, tolerant to the two names the wrapping has shipped under."""
     cls = getattr(itk, "ImpactCoarseRegistration", None) or getattr(itk, "ImpactConvexAdamInitializer", None)
@@ -444,7 +461,7 @@ class ConvexAdamRegistration(torch.nn.Module):
         # ConvexAdam runs a gradient-based instance optimisation (Adam over the field) inside itk-impact's
         # .Update(); the predictor calls forward under torch.inference_mode(), which forbids autograd. The
         # image tensors have already crossed to numpy/ITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad():
+        with torch.inference_mode(False), torch.enable_grad(), _no_texpr_fuser():
             for b in range(fixed.shape[0]):
                 fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
                 moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])

ConvexAdam_Composite/Model.py

@@ -33,6 +33,7 @@ NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engi
 runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
 """
 
+import contextlib
 import json
 import os
 from dataclasses import dataclass, field
@@ -106,6 +107,22 @@ class ModelSpec:
     distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
 
 
+@contextlib.contextmanager
+def _no_texpr_fuser():
+    """Disable ONLY torch's TensorExpr (NNC) fuser for the block. The TS feature-model graphs have a direction/
+    orientation branch whose shape ops (aten::dim / aten::size) crash the fuser's alias analysis under itk-impact
+    ("INTERNAL ASSERT ... We don't have an op for aten::size" in FuseTensorExprs). It cannot be fixed model-side:
+    itk-impact passes an *undefined* direction, so guarding it needs a shape op — the very fuser trigger. The
+    modern profiling executor stays on (NOT the legacy executor); measured cost ~1% of a registration (the fuser
+    only touches the few feature forwards, not the C++ optimisation loop). The caller's setting is restored."""
+    prev = torch._C._jit_texpr_fuser_enabled()
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    try:
+        yield
+    finally:
+        torch._C._jit_set_texpr_fuser_enabled(prev)
+
+
 def _coarse_registration_type():
     """The coupled-convex initializer, tolerant to the two names the wrapping has shipped under."""
     cls = getattr(itk, "ImpactCoarseRegistration", None) or getattr(itk, "ImpactConvexAdamInitializer", None)
@@ -444,7 +461,7 @@ class ConvexAdamRegistration(torch.nn.Module):
         # ConvexAdam runs a gradient-based instance optimisation (Adam over the field) inside itk-impact's
         # .Update(); the predictor calls forward under torch.inference_mode(), which forbids autograd. The
         # image tensors have already crossed to numpy/ITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad():
+        with torch.inference_mode(False), torch.enable_grad(), _no_texpr_fuser():
             for b in range(fixed.shape[0]):
                 fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
                 moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])

ConvexAdam_Fine/Model.py

@@ -33,6 +33,7 @@ NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engi
 runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
 """
 
+import contextlib
 import json
 import os
 from dataclasses import dataclass, field
@@ -106,6 +107,22 @@ class ModelSpec:
     distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
 
 
+@contextlib.contextmanager
+def _no_texpr_fuser():
+    """Disable ONLY torch's TensorExpr (NNC) fuser for the block. The TS feature-model graphs have a direction/
+    orientation branch whose shape ops (aten::dim / aten::size) crash the fuser's alias analysis under itk-impact
+    ("INTERNAL ASSERT ... We don't have an op for aten::size" in FuseTensorExprs). It cannot be fixed model-side:
+    itk-impact passes an *undefined* direction, so guarding it needs a shape op — the very fuser trigger. The
+    modern profiling executor stays on (NOT the legacy executor); measured cost ~1% of a registration (the fuser
+    only touches the few feature forwards, not the C++ optimisation loop). The caller's setting is restored."""
+    prev = torch._C._jit_texpr_fuser_enabled()
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    try:
+        yield
+    finally:
+        torch._C._jit_set_texpr_fuser_enabled(prev)
+
+
 def _coarse_registration_type():
     """The coupled-convex initializer, tolerant to the two names the wrapping has shipped under."""
     cls = getattr(itk, "ImpactCoarseRegistration", None) or getattr(itk, "ImpactConvexAdamInitializer", None)
@@ -444,7 +461,7 @@ class ConvexAdamRegistration(torch.nn.Module):
         # ConvexAdam runs a gradient-based instance optimisation (Adam over the field) inside itk-impact's
         # .Update(); the predictor calls forward under torch.inference_mode(), which forbids autograd. The
         # image tensors have already crossed to numpy/ITK here, so re-enable grad for the optimisation.
-        with torch.inference_mode(False), torch.enable_grad():
+        with torch.inference_mode(False), torch.enable_grad(), _no_texpr_fuser():
             for b in range(fixed.shape[0]):
                 fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
                 moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])