Support C->R case #114
Support C->R case #114
Conversation
…e it behaves differently
|
@patrick-kidger I think this is a reasonable implementation of what we discussed. Some documentation changes are probably required. |
|
ping @patrick-kidger |
| @@ -19,6 +19,7 @@ | |||
| import jax.core | |||
| import jax.numpy as jnp | |||
| import jax.tree_util as jtu | |||
| from jax import ShapeDtypeStruct | |||
There was a problem hiding this comment.
Nit: I'd just use jax.ShapeDtypeStruct directly for consistency with the rest of our imports.
| if is_complex_structure(operator.in_structure()) and not is_complex_structure( | ||
| operator.out_structure() | ||
| ): | ||
| # We'll use R^2->R representation for C->R function. | ||
| in_structure = complex_to_real_structure(operator.in_structure()) | ||
|
|
||
| map_to_original = lambda x: real_to_complex_tree( | ||
| x, | ||
| operator.in_structure(), | ||
| ) | ||
| else: | ||
| map_to_original = lambda x: x | ||
| in_structure = operator.in_structure() | ||
| flat, unravel = strip_weak_dtype( | ||
| eqx.filter_eval_shape(jfu.ravel_pytree, operator.in_structure()) | ||
| eqx.filter_eval_shape(jfu.ravel_pytree, in_structure) | ||
| ) | ||
| fn = lambda x: operator.fn(map_to_original(unravel(x))) | ||
| eye = jnp.eye(flat.size, dtype=flat.dtype) | ||
| jac = jax.vmap(lambda x: operator.fn(unravel(x)), out_axes=-1)(eye) | ||
|
|
||
| jac = jax.vmap(fn, out_axes=-1)(eye) |
There was a problem hiding this comment.
After this change: given a C->R FunctionLinearOperator, call it operator, then I think operator(some_complex_vector) would have worked but materialise(operator)(some_complex_vector) would not work, as it now expects something from R^2 instead?
The intention was that materialise would not change the observable input-output behaviour of an operator at all.
(FWIW I've just checked the current behaviour on FunctionLinearOperator(lambda x: x.real, jax.ShapeDtypeStruct((), jnp.complex64)) and this is also wrong in the expected way: there's no way to express 'take a real part' when multiplying against a pytree, so it's not like the current state of affairs is any better... !)
It seems to me like the AbstractLinearOperator abstraction might actually just be fundamentally incompatible with complex dtypes, due to the not-really-defined nature of linear operators over such spaces?
There was a problem hiding this comment.
You're right, as we discussed in patrick-kidger/optimistix#76 (comment) it is impossible to materialize C->R operator, since it is not complex linear but rather linear in imaginary and real parts. The two solutions I see is one in this PR (break the promise that materialize is noop) or just assert we do not support C->R operators (probably just give a warning) and suggest to user to make R^2->R operator out of it. The latter is definitely cleaner from the point of view of consistency, but may be less convenient to end user? Not sure about it, but luckily for me the decision is yours 🙃
The motivation was to support C->R stuff in optax so it can be a good testbed in terms of understanding how convenient this stuff is
There was a problem hiding this comment.
I think we should follow JAX's lead on this one, with how jax.jac{fwd,rev} similarly error out for these kinds of operations.
So I think something like this might be the best choice, then!
class AbstractLinearOperator(eqx.Module):
def __check_init__(self):
if is_complex_structure(self.in_structure()) and not is_complex_structure(self.out_structure()):
raise ValueError(...)In terms of end user convenience: I care a lot about this! My usual rule for this is that it is better not to support something than to support it awkwardly or with edge-cases. I think this leads to less frustration and a better UX overall. :)
There was a problem hiding this comment.
Fair enough! Note that your solution will break existing cases for mixed-typed operators (i.e., single operator built from C->C and R->R blocks).
There was a problem hiding this comment.
Hmm, that is a good point about block-with-mixed-type operators. Ach, complex support is very complicated! Is there a clean way to detect and allow that case, that you can see?
I am concerned that our complex support isn't quite meeting the above UX standards I'd like us to have... (this has definitely been a learning experience for me on how complex autodiff/operators/etc work!)
There was a problem hiding this comment.
One possible way is to materialize the operator, but that can be expensive. I was thinking just issue a warning instead of exception.
There was a problem hiding this comment.
Hmm. I'm not sure I'm comfortable with the warning approach -- I try to avoid this kind of flaky maybe-right maybe-not behaviour. I'd rather just prohibit C->R altogether, even if it prohibits mixed C->C + R->R combinations. (Which can anway still be trivially supported by an end using by doing C=R^2 themselves.)
Follow up on patrick-kidger/optimistix#76