Optimization of complex parameters
- Dominant language
- Python
- Stars
- 1.1k
- Forks
- 76
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 1
Description
When optimizing a real loss function for a complex variable I encountered some issues with jaxopt 0.7. It seems this use case is less widely used. Are the methods also designed for this use case?
1. I had to use the complex conjugate of the gradient. Is this due to the choice for the complex gradient of jax? Should the jaxopt optimizers take this into account?
2. The 'error' (L2 norm of the gradient), which is used for convergence check is calculated as sum of squares. For complex gradients this should be replaced by the sum over the absolute value, or equivalently, the sum of the squared real and imaginary parts. As a simple fix a patched the l-BGFS and GradientDescent methods. I tested that this change also works with real gradients, and when jitted the extra operation is removed for real gradients.
```python
def tree_l2_norm_patched(tree_x, squared=False):
"""Compute the l2 norm ||tree_x||."""
squared_tree = jaxopt.tree_util.tree_map(lambda x: jnp.square(x.real) + jnp.square(x.imag), tree_x)
sqnorm = jaxopt.tree_util.tree_sum(squared_tree)
if squared:
return sqnorm
else:
return jnp.sqrt(sqnorm)
import jaxopt._src
jaxopt._src.lbfgs.tree_l2_norm = tree_l2_norm_patched
jaxopt._src.proximal_gradient.tree_l2_norm = tree_l2_norm_patched
```
Gregor
Contributor guide
Assessment
This issue has not been assessed yet.