tf.control_dependencies

tf.control_dependencies(
    control_inputs
) -> Graph._ControlDependenciesController

@tf.function(
  input_signature=[tf.TensorSpec([None, None], tf.float32),
                   tf.TensorSpec([None, None], tf.float32)])
def my_assert_func_1(x, bias):
  # `tf.function` attempts to execute `tf.math.add` in parallel to
  # `assert_equal`. As a result an error can get raised from `tf.math.add`
  # without triggering the assertion error.
  tf.assert_equal(tf.shape(x)[1],
                  tf.shape(bias)[1],
                  message='bad shape')
  return x + bias

# Error raised in either `add` or `assert`
my_assert_func_1(tf.ones((2, 5)), tf.ones((2, 7)))
Traceback (most recent call last):

InvalidArgumentError: ...

@tf.function(
  input_signature=[tf.TensorSpec([None, None], tf.float32),
                   tf.TensorSpec([None, None], tf.float32)])
def my_assert_func_2(x, bias):
  with tf.control_dependencies(
      [tf.assert_equal(tf.shape(x)[1],
                      tf.shape(bias)[1],
                      message='bad shape')]):
    return x + bias

# Error raised in `assert`
my_assert_func_2(tf.ones((2, 5)), tf.ones((2, 7)))
Traceback (most recent call last):

InvalidArgumentError: ...