Skip to content
8 changes: 4 additions & 4 deletions pymc_marketing/mmm/mmm.py
Original file line number Diff line number Diff line change
Expand Up @@ -4029,7 +4029,7 @@ def optimize_budget(
budget_distribution_over_period: xr.DataArray | None = None,
cost_per_unit: pd.DataFrame | xr.DataArray | None = None,
callback: bool = False,
**minimize_kwargs,
**allocate_budget_kwargs,
) -> (
tuple[xr.DataArray, OptimizeResult]
| tuple[xr.DataArray, OptimizeResult, list[dict[str, Any]]]
Expand Down Expand Up @@ -4080,8 +4080,8 @@ def optimize_budget(
**This is independent of the historical cost_per_unit.**
callback : bool
Whether to return callback information tracking optimization progress.
**minimize_kwargs
Additional arguments for the optimizer.
**allocate_budget_kwargs
Additional arguments for :meth:`~pymc_marketing.mmm.budget_optimizer.BudgetOptimizer.allocate_budget`.

Returns
-------
Expand Down Expand Up @@ -4130,7 +4130,7 @@ def optimize_budget(
total_budget=budget,
budget_bounds=budget_bounds,
callback=callback,
**minimize_kwargs,
**allocate_budget_kwargs,
)

def _apply_budget_distribution_pattern(
Expand Down
Loading