evanatyourservice

import json
from tqdm import tqdm
import numpy as np

import jax
import jax.numpy as jnp
from flax.training.train_state import TrainState
import optax
import optax.tree_utils as otu
import tensorflow as tf

evanatyourservice

import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from tensorflow_probability.substrates.jax import distributions as tfd


"""
There's a typo in most B-TCVAE implementations on github, so I thought I'd make a 
quick gist of a working B-TCVAE.