python import shorthands
from tensorflow_probability.substrates import jax as tfp
import tensorflow_probability.substrates.jax as tfp