tensorflow_probability: bfgs_minimize requires the elements in initial_position to have the same shape

This line:
https://github.com/tensorflow/probability/blob/main/tensorflow_probability/python/optimizer/bfgs.py#L210
fails for input of the form:
initial_position = [0., tf.zeros(10)]

I am working through some simple toy examples and have a JointDistribution of the form:

def make_joint_dist(matrix):
    def joint_dist():
        intercept = yield tfd.Normal(loc=0.0, scale=1.0, name="intercept")
        coefficients = yield tfd.Normal(loc=tf.zeros(10), scale=1.0, name="coefficients")
        yield tfd.Normal(
            loc=intercept + tfl.matvec(matrix, coefficients),
            scale=1.0,
            name="observations",
        )

    return tfd.JointDistributionCoroutineAutoBatched(joint_dist)

I am trying to run bfgs_minimize to find the MAP (minimize -joint_dist.log_prob(params)). However, in this case, I am cannot figure out how to pass an input_position for the coefficients. The code fails at L210 above:
Shapes of all inputs must match: values[0].shape = [] != values[1].shape = [10]

I can put together a running example colab, but wanted to kick this off while I clean up my notebook....

Start at tensorflow_probability/python/optimizer/bfgs.py line 210 and reproduce the failure with initial_position = [0., tf.zeros(10)]. Trace how bfgs_minimize combines the initial values and determine the expected behavior for differently shaped elements. Done means this input no longer raises the reported shape-mismatch error.