jax.grad of a real-valued estimator returns the conjugated Wirtinger derivative for complex-valued parameters. For example, where the true gradient is $(\partial f/\partial x, \partial f/\partial y) = (4, 7)$, jax.grad returns 4-7j. ParameterFlattener.flatten() maps imag_<name> to value.imag, so optimizers receive $-\partial f/\partial y$ (wrong sign) for every complex parameter when fitting with use_analytic_gradient=True.
Fix: conjugate the complex leaves of the jax.grad output in gradient_creator(). The existing gradient tests only cover real-valued parameters, so add a complex-parameter case with a known analytic gradient.
jax.gradof a real-valued estimator returns the conjugated Wirtinger derivative for complex-valued parameters. For example, where the true gradient isjax.gradreturns4-7j.ParameterFlattener.flatten()mapsimag_<name>tovalue.imag, so optimizers receiveuse_analytic_gradient=True.Fix: conjugate the complex leaves of the
jax.gradoutput ingradient_creator(). The existing gradient tests only cover real-valued parameters, so add a complex-parameter case with a known analytic gradient.