def log_prob(self, value):
if self._validate_args:
self._validate_sample(value)
dtype = get_dtypes(self.logits)[0]
value = lax.convert_element_type(value, dtype)
return -binary_cross_entropy_with_logits(self.logits, value)
def test_binary_cross_entropy_with_logits(x, y):
actual = -y * jnp.log(expit(x)) - (1 - y) * jnp.log(expit(-x))
expect = binary_cross_entropy_with_logits(x, y)
assert_allclose(actual, expect, rtol=1e-6)
def _logpmf(self, x, p):
if self.is_logits:
return -binary_cross_entropy_with_logits(p, x)
else:
# TODO: consider always clamp and convert probs to logits
return xlogy(x, p) + xlog1py(1 - x, -p)
def log_prob(self, value):
return -binary_cross_entropy_with_logits(self.logits, value)
def log_prob(self, value):
if self._validate_args:
self._validate_sample(value)
return -binary_cross_entropy_with_logits(self.logits, value)
