evanatyourservice

import json
from typing import Optional
from tqdm import tqdm
import numpy as np

import jax
from jax import jit
import jax.numpy as jnp
from flax.training.train_state import TrainState
import optax

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.