| 733 | |
| 734 | |
| 735 | class FlaxDiagonalGaussianDistribution(object): |
| 736 | def __init__(self, parameters, deterministic=False): |
| 737 | # Last axis to account for channels-last |
| 738 | self.mean, self.logvar = jnp.split(parameters, 2, axis=-1) |
| 739 | self.logvar = jnp.clip(self.logvar, -30.0, 20.0) |
| 740 | self.deterministic = deterministic |
| 741 | self.std = jnp.exp(0.5 * self.logvar) |
| 742 | self.var = jnp.exp(self.logvar) |
| 743 | if self.deterministic: |
| 744 | self.var = self.std = jnp.zeros_like(self.mean) |
| 745 | |
| 746 | def sample(self, key): |
| 747 | return self.mean + self.std * jax.random.normal(key, self.mean.shape) |
| 748 | |
| 749 | def kl(self, other=None): |
| 750 | if self.deterministic: |
| 751 | return jnp.array([0.0]) |
| 752 | |
| 753 | if other is None: |
| 754 | return 0.5 * jnp.sum(self.mean**2 + self.var - 1.0 - self.logvar, axis=[1, 2, 3]) |
| 755 | |
| 756 | return 0.5 * jnp.sum( |
| 757 | jnp.square(self.mean - other.mean) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar, |
| 758 | axis=[1, 2, 3], |
| 759 | ) |
| 760 | |
| 761 | def nll(self, sample, axis=[1, 2, 3]): |
| 762 | if self.deterministic: |
| 763 | return jnp.array([0.0]) |
| 764 | |
| 765 | logtwopi = jnp.log(2.0 * jnp.pi) |
| 766 | return 0.5 * jnp.sum(logtwopi + self.logvar + jnp.square(sample - self.mean) / self.var, axis=axis) |
| 767 | |
| 768 | def mode(self): |
| 769 | return self.mean |
| 770 | |
| 771 | |
| 772 | @flax_register_to_config |
no outgoing calls
no test coverage detected
searching dependent graphs…