Add Exp and Softplus#67
Open
gvcallen wants to merge 3 commits into
Open
Conversation
lockwo
reviewed
Apr 6, 2026
| """ | ||
|
|
||
| _is_constant_jacobian: bool = True | ||
| _is_constant_log_det: bool = True |
Owner
There was a problem hiding this comment.
These depend on x right, so should be False?
| self.assertion_fn()(log_det, -jnp.log(y)) | ||
| self.assertEqual(x.dtype, dtype) | ||
| self.assertEqual(log_det.dtype, dtype) | ||
|
|
Owner
There was a problem hiding this comment.
I like these round trip tests, e.g. the inverse(forward(x)) == x style, can we add one for exp (since there is one for soft plus)?
gvcallen
added a commit
to gvcallen/distreqx
that referenced
this pull request
Jul 1, 2026
Softplus's log-det depends on x, so _is_constant_jacobian/_is_constant_log_det must be False (matching lockwo's observation), not True. Add a forward/inverse round-trip test to exp_test.py matching the existing softplus test pattern.
Exponential and softplus bijectors modeled after TensorFlow Probability's implementations, mapping the real line to the positive domain.
gvcallen
force-pushed
the
exp_and_softplus
branch
from
89d7d61 to
98b5be0
Compare
_more_stable_sigmoid and _more_stable_softplus used jnp.where(cond, a, b) to switch to a numerically stable approximation for very negative inputs. jnp.where evaluates both branches unconditionally, so for large positive inputs the unselected branch (jnp.exp(x) / log1p(exp(x))) overflowed to inf, and the 0 * inf this produces in the backward pass poisoned the gradient with NaN even though that branch was never used for the forward value. Fix by feeding the unselected branch a safe input via a second jnp.where, the standard remedy for this "double where" trap. Forward values are unchanged; only the previously-poisoned gradient at large positive inputs is fixed. Adds a regression test exercising both helpers (via Sigmoid.forward and Sigmoid.forward_log_det_jacobian) across a range including extreme positive and negative inputs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds the exponential and softplus bijectors like Tensorflow's tfp.bijectors.Exp and tfp.bijectors.Softplus.