GPU support in BudgetOptimizer (#1937)

* Introduce jax_backend boolean to enable jax-based optimizer computation

* Jaxify _compiled_functions

* Add extensive jax_backend tests

* BudgetOptimizer now accepts compile_kwargs dict

* Adjust wrapper to BudgetOptimizer signature

* Update tests

Author: PabloRoque

pymc_marketing/mmm/budget_optimizer.py

@@ -363,6 +363,11 @@ class BudgetOptimizer(BaseModel):
         ),
     )
 
+    compile_kwargs: dict | None = Field(
+        default=None,
+        description="Keyword arguments for the model compilation. Specially usefull to pass compilation mode",
+    )
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     DEFAULT_MINIMIZE_KWARGS: ClassVar[dict] = {
@@ -684,10 +689,24 @@ def _compile_objective_and_grad(self):
         )
         objective_grad = pt.grad(rewrite_pregrad(objective), budgets_flat)
 
-        objective_and_grad_func = function([budgets_flat], [objective, objective_grad])
+        if self.compile_kwargs and (self.compile_kwargs["mode"]).lower() == "jax":
+            # Use PyMC's JAX infrastructure for robust compilation
+            from pymc.sampling.jax import get_jaxified_graph
+
+            objective_and_grad_func = get_jaxified_graph(
+                inputs=[budgets_flat],
+                outputs=[objective, objective_grad],
+                **{k: v for k, v in self.compile_kwargs.items() if k != "mode"} or {},
+            )
+        else:
+            # Standard PyTensor compilation
+            objective_and_grad_func = function(
+                [budgets_flat], [objective, objective_grad], **self.compile_kwargs or {}
+            )
 
         # Avoid repeated input validation for performance
-        objective_and_grad_func.trust_input = True
+        if hasattr(objective_and_grad_func, "trust_input"):
+            objective_and_grad_func.trust_input = True
 
         self._compiled_functions[self.utility_function] = {
             "objective_and_grad": objective_and_grad_func,

pymc_marketing/mmm/multidimensional.py

@@ -2146,7 +2146,13 @@ def create_sample_kwargs(
 class MultiDimensionalBudgetOptimizerWrapper(OptimizerCompatibleModelWrapper):
     """Wrapper for the BudgetOptimizer to handle multi-dimensional model."""
 
-    def __init__(self, model: MMM, start_date: str, end_date: str):
+    def __init__(
+        self,
+        model: MMM,
+        start_date: str,
+        end_date: str,
+        compile_kwargs: dict | None = None,
+    ):
         self.model_class = model
         self.start_date = start_date
         self.end_date = end_date
@@ -2158,6 +2164,7 @@ def __init__(self, model: MMM, start_date: str, end_date: str):
             include_carryover=False,
         )
         self.num_periods = len(self.zero_data[self.model_class.date_column].unique())
+        self.compile_kwargs = compile_kwargs
         # Adding missing dependencies for compatibility with BudgetOptimizer
         self._channel_scales = 1.0
 
@@ -2256,6 +2263,7 @@ def optimize_budget(
             budgets_to_optimize=budgets_to_optimize,
             budget_distribution_over_period=budget_distribution_over_period,
             model=self,  # Pass the wrapper instance itself to the BudgetOptimizer
+            compile_kwargs=self.compile_kwargs,
         )
 
         return allocator.allocate_budget(