Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add conj function and adjoint tests #62

Merged
merged 4 commits into from
Nov 12, 2023

Conversation

Randl
Copy link
Contributor

@Randl Randl commented Nov 1, 2023

Add adjoint property, which is useful for complex-valued operators.
See #57

@patrick-kidger
Copy link
Owner

What's the use-case for adding this? The only reason we have a transpose function is to handle the transposition necessary to synthesise the backward (VJP) pass. In particular JAX performs only transposition (in contrast to PyTorch, which performs conjugate-transposition), so I don't think we actually need this for any internal operation.

@Randl
Copy link
Contributor Author

Randl commented Nov 1, 2023

Specifically for normal operators, e.g., in CG part

@patrick-kidger
Copy link
Owner

Hmm. I think we can probably handle that without adding a new API. Rewriting this line:
https://github.com/google/lineax/blob/9a573529f3eea2a6944009f6e4afe4f364366baf/lineax/_solver/cg.py#L120
to be ω(_transpose_mv(ω(_mv(vector)).call(jnp.conj).ω)).call(jnp.conj).ω should work I think. (This is morally just a jnp.conj(_transpose_mv(jnp.conj(_mv(vector)))), the ω is just a neat syntax we have for tree-mapping an operation. And check my math but I think conj(A conj(b)) = conj(A) b.)

Mostly I'm trying to avoid adding more to the public API. We'd also then need to register rules for every downstream operator.

@Randl
Copy link
Contributor Author

Randl commented Nov 1, 2023

I see the reasoning, but currently it's suggested that the user should do the adjoint themselves at least in some cases:

https://github.com/google/lineax/blob/9a573529f3eea2a6944009f6e4afe4f364366baf/lineax/_solver/cg.py#L107-L112

Also, not having adjoint means bilinears are out of the question, if this is relevant.

@patrick-kidger
Copy link
Owner

Hmm.
Okay, I think I buy your argument. In that case I think I'd suggest adding a lineax.conj function, using single-dispatch in the same way as lineax.{linearise, ...}. This means that we'll be compatible with any other existing AbstractLinearOperator subclasses out there (no new abstract method), and also means we don't need to worry about the complexities of doing the full adjoint in one go -- we can decompose into conjugation and transposition, which should be simpler to reason about.

@Randl Randl changed the title Add adjoint function and tests Add conj function and adjoint tests Nov 2, 2023
Copy link
Owner

@patrick-kidger patrick-kidger left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nits aside, this LGTM! Thank you for taking the time to make this happen.

lineax/__init__.py Outdated Show resolved Hide resolved
lineax/_operator.py Outdated Show resolved Hide resolved
lineax/_operator.py Outdated Show resolved Hide resolved
lineax/_operator.py Outdated Show resolved Hide resolved
@patrick-kidger
Copy link
Owner

By the way, let me know what your plan is regarding the various PRs. I think what might be easiest is to finish all of them before merging any of them, and then we can merge them in whatever order you think is best.

(Some good news - I have no changes planned for Lineax at the moment, so you shouldn't need to worry about keeping any of them up to date with any other changes.)

@Randl
Copy link
Contributor Author

Randl commented Nov 3, 2023

From my perspective, these 5 PRs are ready to merge (up to minor unclosed discussions). It does not fully implement complex support, but it would be simpler for me to continue when these are merged so I can be sure this part already works.

As for the merge order, I think it should #59 and #63 should go first, then #60, then #61 and this one

@patrick-kidger patrick-kidger merged commit 49e83ef into patrick-kidger:main Nov 12, 2023
3 checks passed
@patrick-kidger
Copy link
Owner

Also LGTM!

@Randl Randl mentioned this pull request Nov 12, 2023
8 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants