Jax ml operator fix

firedrake/ml/jax/ml_operator.py

@@ -90,7 +90,10 @@ def _pre_forward_callback(self, *operands: Union[Function, Cofunction], unsqueez
 
     def _post_forward_callback(self, y_P: "jax.Array") -> Union[Function, Cofunction]:
         """Callback function to convert the JAX output of the ML model to a Firedrake function."""
-        space = self.ufl_function_space()
+        if isinstance(self._argument_slots[0], Cofunction):
+            space = self.ufl_operands[0].function_space().dual()
+        else:
+            space = self.ufl_function_space()
         return from_jax(y_P, space)
 
     # -- JAX routines for computing AD-based quantities -- #

firedrake/ml/pytorch/ml_operator.py

@@ -98,7 +98,11 @@ def _pre_forward_callback(self, *operands, unsqueeze=False):
 
     def _post_forward_callback(self, y_P):
         """Callback function to convert the PyTorch output of the ML model to a Firedrake function."""
-        space = self.ufl_function_space()
+        from firedrake import Cofunction
+        if isinstance(self._argument_slots[0], Cofunction):
+            space = self.ufl_function_space().dual()
+        else:
+            space = self.ufl_function_space()
         return from_torch(y_P, space)
 
     # -- PyTorch routines for computing AD based quantities via `torch.autograd.functional` -- #