Feasibility polishing doesn't work

[Richienb]

The error I get:

File ~/<my project>/.venv/lib/python3.12/site-packages/mpax/feasibility_polishing.py:13, in init_primal_feasibility_polishing(scaled_problem, solver_state, initial_primal_weight, average)
     10 zero_objective_matrix = jnp.zeros_like(scaled_problem.original_qp.objective_matrix)
     11 zero_objective_vector = jnp.zeros_like(scaled_problem.original_qp.objective_vector)
     12 primal_feasibility_problem = ScaledQpProblem(
---> 13     original_qp=scaled_problem.original_qp.replace(
     14         objective_matrix=zero_objective_matrix,
     15         objective_vector=zero_objective_vector,
     16         objective_constant=0,
     17     ),
     18     scaled_qp=scaled_problem.scaled_qp.replace(
     19         objective_matrix=zero_objective_matrix,
     20         objective_vector=zero_objective_vector,
     21         objective_constant=0,
     22     ),
     23     constraint_rescaling=scaled_problem.constraint_rescaling,
     24     variable_rescaling=scaled_problem.variable_rescaling,
     25 )
     26 zero_dual_solution = jnp.zeros_like(solver_state.current_dual_solution)
     27 zero_dual_product = jnp.zeros_like(solver_state.current_dual_product)

AttributeError: 'QuadraticProgrammingProblem' object has no attribute 'replace'

It's referring to this code:

MPAX/mpax/feasibility_polishing.py

Lines 12 to 22 in ca1c669

	primal_feasibility_problem = ScaledQpProblem(
	original_qp=scaled_problem.original_qp.replace(
	objective_matrix=zero_objective_matrix,
	objective_vector=zero_objective_vector,
	objective_constant=0,
	),
	scaled_qp=scaled_problem.scaled_qp.replace(
	objective_matrix=zero_objective_matrix,
	objective_vector=zero_objective_vector,
	objective_constant=0,
	),

It looks like you can make .replace() accessible by using the chex.dataclass decorator, rather than the regular decorator here:

MPAX/mpax/r2hpdhg.py

Lines 49 to 53 in ca1c669

	@dataclass(eq=False)
	class r2HPDHG(raPDHG):
	"""
	The r2HPDHG solver class.
	"""

Alternatively, you can fix the calling code by using dataclasses.replace, rather than calling .replace().

[ZedongPeng]

Thank you @Richienb . I will fix this.