ffjord bijector takes so long to process one mnist size data with a simple ODEnet

Here's the code. I use the code from official ffjord demo but instead of using snt I use tf.keras.

base_loc = np.array([0.0, 0.0]).astype(np.float32)
base_sigma = np.array([0.8, 0.8]).astype(np.float32)
base_distribution = tfd.MultivariateNormalDiag(base_loc, base_sigma)

class MLP_ODE(tf.keras.Model):
  """Multi-layer NN ode_fn."""
  def __init__(self, num_hidden, num_layers, num_output, name='mlp_ode'):
    super(MLP_ODE, self).__init__(name=name)
    self._num_hidden = num_hidden
    self._num_output = num_output
    self._num_layers = num_layers
    self._modules = []
    for _ in range(self._num_layers - 1):
      self._modules.append(tf.keras.layers.Dense(self._num_hidden, activation='tanh'))
    self._modules.append(tf.keras.layers.Dense(self._num_output))
    self._model = tf.keras.Sequential(self._modules)

  def __call__(self, t, inputs):
    inputs = tf.concat([tf.broadcast_to(t, inputs.shape), inputs], -1)
    return self._model(inputs)

#@title Model and training parameters
LR = 1e-2  #@param
NUM_EPOCHS = 80  #@param
STACKED_FFJORDS =   4#@param
NUM_HIDDEN = 8  #@param
NUM_LAYERS = 3  #@param
NUM_OUTPUT = 2

solver = tfp.math.ode.DormandPrince(atol=1e-5)
ode_solve_fn = solver.solve
trace_augmentation_fn = tfb.ffjord.trace_jacobian_exact

bijectors = []
for _ in range(STACKED_FFJORDS):
  mlp_model = MLP_ODE(NUM_HIDDEN, NUM_LAYERS, NUM_OUTPUT)
  next_ffjord = tfb.FFJORD(
      state_time_derivative_fn=mlp_model,ode_solve_fn=ode_solve_fn,
      trace_augmentation_fn=trace_augmentation_fn)
  bijectors.append(next_ffjord)

stacked_ffjord = tfb.Chain(bijectors[::-1])

transformed_distribution = tfd.TransformedDistribution(
    distribution=base_distribution, bijector=stacked_ffjord)
test_val = tf.Variable(np.ones((32,32,2)), dtype=tf.float32)

⚠️⚠️⬇️This line of code takes almost 2min on colab with gpu support.⬇️

print(transformed_distribution.log_prob(test_val))

Why it takes so long to calculate the log_prob? Am I using this bijector in a wrong way?

Start at transformed_distribution.log_prob(test_val) in the provided notebook code and profile the FFJORD chain, DormandPrince solver, and trace_jacobian_exact configuration. Compare the runtime with the shown batch shape and determine whether the delay comes from the solver or the bijector usage; done means explaining the bottleneck and confirming the appropriate usage.