Skip to content

Add Inverse and Identity#64

Open
gvcallen wants to merge 1 commit intolockwo:mainfrom
gvcallen:identity_and_inverse
Open

Add Inverse and Identity#64
gvcallen wants to merge 1 commit intolockwo:mainfrom
gvcallen:identity_and_inverse

Conversation

@gvcallen
Copy link
Copy Markdown

@gvcallen gvcallen commented Apr 4, 2026

Implements the simple utility Inverse and Identity bijectors.

@gvcallen gvcallen changed the title Implement Inverse and Identity Add Inverse and Identity Apr 4, 2026
@gvcallen gvcallen changed the title Add Inverse and Identity Add Inverse and Identity bijectors Apr 4, 2026
@gvcallen gvcallen changed the title Add Inverse and Identity bijectors Add Inverse and Identity Apr 4, 2026

def forward_and_log_det(self, x: PyTree) -> tuple[PyTree, PyTree]:
"""Computes y = x and log|det J(f)(x)| = 0."""
log_det = jax.tree_util.tree_map(jnp.zeros_like, x)
Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

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

Can just be jax.tree.map?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Oh yes, would you like me to commit changes like these?

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

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

should be tested with more complicated bijector than just identity

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