<jupyter_start><jupyter_text><jupyter_code>!git clone https://github.com/tfrizza/DALL-E-tf.git
%cd DALL-E-tf
import tensorflow as tf
from tensorflow.keras import Model, mixed_precision

import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
from tensorflow_probability import layers as tfpl

from dall_e_tf.encoder import dvae_encoder
from dall_e_tf.decoder import dvae_decoder
from dall_e_tf.utils import plot_reconstructions

mixed_precision.set_global_policy('mixed_float16')

(x_train, _), (x_test, _) = tf.keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_train = tf.expand_dims(x_train, -1)
x_test = tf.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
vocab_size = 32
n_hid = 32

enc = dvae_encoder(img_size=28, group_count=3, n_hid=n_hid, n_blk_per_group=2, input_channels=x_train.shape[-1], vocab_size=vocab_size)
dec = dvae_decoder(img_size=7, group_count=3, n_init=n_hid//2, n_hid=n_hid, n_blk_per_group=2, output_channels=x_train.shape[-1], vocab_size=vocab_size)

def dVAE(enc, dec, temperature=1/16):
    z_e = enc.output # encoder logits

    z_q = tfpl.DistributionLambda(
        lambda logits: tfd.RelaxedOneHotCategorical(temperature, logits) # Gumbel-softmax
    )(z_e)

    z_hard = tf.math.argmax(z_e, axis=-1) # non-differentiable
    z_hard = tf.one_hot(z_hard, enc.output.shape[-1], dtype=tf.float16)

    z = z_q + tf.stop_gradient(z_hard - z_q) # straight-through Gumbel-softmax
    x_rec = dec(z)
    return Model(enc.input, x_rec)

vae = dVAE(enc, dec, 0.1)
vae.summary(line_length=200)
vae.compile(loss='mse', optimizer='adam')

vae.fit(x_train, x_train,
        validation_data=(x_test, x_test),
        epochs=1,
        batch_size=32,
        shuffle=True)
plot_reconstructions(vae(x_test[:10]), x_test[:10])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12457254518808011, 'fraction_numerical': 0.027357107962872496, 'mean_word_length': 5.340557275541796, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '197806', 'license': 'MIT', 'n_tokens_mistral': 858, 'n_tokens_neox': 815, 'n_words': 152, 'path': 'MNIST_demo.ipynb', 'repo_name': 'tfrizza/DALL-E-tf', 'size': '2047'}