Add Exp and Softplus - #67
Conversation
| """ | ||
|
|
||
| _is_constant_jacobian: bool = True | ||
| _is_constant_log_det: bool = True |
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) | ||
|
|
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)?
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.
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.
|
btw, just in general, when the comments are changed/fixed you can just resolved them, then when it is ready for another round of review, you can click the request review again button. I have a lot of GitHub notifications, so it isn't as likely that I will know it is ready without this |
|
@lockwo okay thanks! Sorry I am quite new to open-source workflows! |
| def _more_stable_sigmoid(x: Array) -> Array: | ||
| """Where extremely negatively saturated, approximate sigmoid with exp(x).""" | ||
| ret = jnp.where(x < -9, jnp.exp(x), jax.nn.sigmoid(x)) | ||
| # `jnp.where` evaluates both branches, so the unselected `jnp.exp(x)` branch |
There was a problem hiding this comment.
this is an interesting issue, I copied this code from the original distrax so I'm surprised they didn't run into this issue. Did you encounter this in your work?
There was a problem hiding this comment.
@lockwo yes I did actually! I can't remember where though, I think it was during an optimization problem using the bijector for constraints, but I definitely remember bumping into it and having to fix it
| ): | ||
| """ | ||
| Transforms the real line to the positive domain using | ||
| softplus y = log(1 + exp(x)). |
| @@ -0,0 +1,40 @@ | |||
| import jax.nn as jnn | |||
There was a problem hiding this comment.
I prefer just import jax then later do jax.nn
2 participants
Adds the exponential and softplus bijectors like Tensorflow's tfp.bijectors.Exp and tfp.bijectors.Softplus.