TFP optimizers should take an "additional_args" argument which is passed through to value_and_gradients_function

Consider the following scenario (shamelessly copied from TFP tutorial)

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
np.random.seed(12345)

@tf.function
def quadratic(minimum, x):
  with tf.GradientTape() as g:
      g.watch(x)
      out = tf.reduce_sum(input_tensor=scales * (x - minimum)**2, axis=-1)
  grad = g.gradient(out, x)
  return out, grad

We wish to minimize quadratic with respect to a specific argument minimum. However, tfp.optimizers does not allow passing any additional arguments "straight through" to the value_and_gradients function.

This is a problem as it forces eager execution like as follows:

def quadratic_wrapper(minimum):
    return (lambda x: quadratic(minimum, x))

Now quadratic wrapper cannot be turned into a tf.function as it returns a closure, not a tensor. So we must do something like this:

tfp.optimizer.lbfgs_minimize(quadratic_wrapper(minimum), ...

Now the above also cannot be turned into a tf.function as it takes a python function as an argument.

Basically, we're forced into eager execution, when, if TFP allowed to pass arguments to value_and_gradients_function in a "passthrough" manner, we could overcome the above as "minimum" in the above example can be represented as a tensor.

I've hacked together a fix for this, and I notice a 2x speedup in my code by implementing something like this. Would be nice if this was officially supported.

The issue names tfp.optimizers and value_and_gradients_function but no files or tests; start by locating the optimizer implementations and their existing call paths. Confirm how an additional_args parameter should be passed through and verify that the tutorial-style optimization works under tf.function without a closure.