File size: 2,813 Bytes
74e8f2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
# Copyright 2024 Big Vision Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Abstract VAE model class.
Gaussian encoder and decoder (the latter assumed to have constant variance).
Inspiration drawn from https://github.com/pytorch/examples/tree/main/vae.
"""
import abc
from typing import Optional, Mapping
import flax.linen as nn
import jax
import jax.numpy as jnp
class Model(nn.Module, metaclass=abc.ABCMeta):
"""Abstract VAE model class."""
codeword_dim: Optional[int] = None
code_len: int = 256
code_dropout: str = "none"
@abc.abstractmethod
def encode(
self,
x: jax.Array,
*,
train: bool = False,
) -> tuple[jax.Array, jax.Array]:
...
def reparametrize(
self,
mu: jax.Array,
logvar: jax.Array,
rng: jax.Array | None = None,
) -> jax.Array:
std = jnp.exp(0.5 * logvar)
if rng is None:
rng = self.make_rng("dropout")
eps = jax.random.normal(rng, shape=std.shape, dtype=std.dtype)
return mu + std * eps
@abc.abstractmethod
def decode(
self, x: jax.Array,
train: bool = False,
) -> jax.Array | Mapping[str, jax.Array]:
...
def code_dropout_fn(self, z: jax.Array, *, train: bool = False) -> jax.Array:
# "seq" drops out tokens later in the sequence with higher probablility than
# tokens earlier in the sequence.
assert self.code_dropout in ["none", "seq", "random"]
if train and self.code_dropout != "none":
importance = jnp.linspace(1.0, 0.0, self.code_len + 2)[1:-1]
thr = jax.random.uniform(self.make_rng("dropout"), z.shape[:1])
mask = importance[None, :] > thr[:, None]
if self.code_dropout == "random":
mask = jax.random.permutation(
self.make_rng("dropout"), mask, axis=-1, independent=True)
z = z * mask[:, :, None]
return z
def __call__(
self,
x: jax.Array,
*,
train: bool = False,
) -> tuple[jax.Array | Mapping[str, jax.Array], Mapping[str, jax.Array]]:
mu, logvar = self.encode(x, train=train)
# Only reparametrize when training for simplicity.
if train:
z = self.reparametrize(mu, logvar)
else:
z = mu
z = self.code_dropout_fn(z, train=train)
x = self.decode(z, train=train)
return x, {"mu": mu, "logvar": logvar, "z": z}
|