LBFGS not working with tensorflow 2.0

System information

• OS Platform and Distribution: Linux NixOS unstable
• TensorFlow installed from : binary using anaconda
• TensorFlow version : '2.0.0-alpha0'
• TensorFlow Probability version : '0.7.0-dev20190504'
• Python version: 3.6.8

Describe the current behavior

When trying to use the tfp.optimizer.lbfgs_minimize function, I get an error, :

InvalidArgumentError: Inputs to operation Select of type Select must have the same size and shape.  Input 0: [1,2] != input 1: [] [Op:Select]

Describe the expected behavior

This should run without issue, as it works under TensorFlow 1.13.1 and TensorFlow Probability 0.6.0, with tf.enable_eager_execution()

Code to reproduce the issue
The following code runs under TF 1.13.1 with TFP 0.6.0, but not with TF 2.0 with TFP 0.7.0

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

class TestEager():
    def __init__(self):
        # tf.losses.mean_squarred_error is not the same under TF 2.0
        self.mse = tf.losses.mean_squared_error
        if tf.__version__ == '2.0.0-alpha0':
            self.mse = tf.losses.MeanSquaredError()

    def __call__(self, inputs):
        loss = 0
        with tf.GradientTape() as tape:
            tape.watch(inputs)
            new_guess = np.random.rand(*inputs.shape)
            loss += self.mse(inputs, new_guess) 
        grad = tape.gradient(loss, inputs)
        return loss, grad

def main_eager():
    guess = np.random.rand(1,2,3).astype(np.float32)
    test_eager = TestEager()
    res = tfp.optimizer.lbfgs_minimize(
      test_eager, 
      initial_position=guess,
      tolerance=1e-8)
    print(res)

if __name__ == "__main__":
    version = tf.__version__
    if version == '2.0.0-alpha0':
        main_eager()
    else:
        tf.enable_eager_execution()
        main_eager()

Other info / logs
Traceback :

2019-05-08 14:58:50.115047: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-05-08 14:58:50.147511: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 1800000000 Hz
2019-05-08 14:58:50.148229: I tensorflow/compiler/xla/service/service.cc:162] XLA service 0x55a747566a30 executing computations on platform Host. Devices:
2019-05-08 14:58:50.148279: I tensorflow/compiler/xla/service/service.cc:169]   StreamExecutor device (0): <undefined>, <undefined>
Traceback (most recent call last):
  File "reprodcuing_bug.py", line 176, in <module>
    main_eager()
  File "reprodcuing_bug.py", line 159, in main_eager
    tolerance=1e-8)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/lbfgs.py", line 260, in minimize
    parallel_iterations=parallel_iterations)[0]
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3216, in while_loop_v2
    return_same_structure=True)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3442, in while_loop
    loop_vars = body(*loop_vars)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/lbfgs.py", line 238, in _body
    tolerance, f_relative_tolerance, x_tolerance, stopping_condition)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/bfgs_utils.py", line 153, in line_search_step
    converged=inactive)  # No search needed for these.
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/linesearch/hager_zhang.py", line 283, in hager_zhang
    right=hzl.val_where(init_converged, val_0, val_c))
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/linesearch/internal/hager_zhang_lib.py", line 45, in val_where
    return cls(*(val_where(cond, t, f) for t, f in zip(tval, fval)))
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/linesearch/internal/hager_zhang_lib.py", line 45, in <genexpr>
    return cls(*(val_where(cond, t, f) for t, f in zip(tval, fval)))
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow_probability/python/optimizer/linesearch/internal/hager_zhang_lib.py", line 42, in val_where
    return tf.where(cond, tval, fval)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow/python/ops/array_ops.py", line 3231, in where
    return gen_math_ops.select(condition=condition, x=x, y=y, name=name)
  File "/home/beren/.conda/envs/style_transfer/lib/python3.6/site-packages/tensorflow/python/ops/gen_math_ops.py", line 9060, in select
    _six.raise_from(_core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: Inputs to operation Select of type Select must have the same size and shape.  Input 0: [1,2] != input 1: [] [Op:Select]

Start by running the supplied reproducer with TensorFlow 2.0.0-alpha0 and TensorFlow Probability 0.7.0-dev20190504. Read tensorflow_probability/python/optimizer/lbfgs.py, bfgs_utils.py, and optimizer/linesearch/internal/hager_zhang_lib.py around the failing val_where call; done means lbfgs_minimize runs for the shown eager-mode example without the Select shape error.