tensorflow - tf.cond does not work if predicate is passed from function
@eml-39502,
Looks like this is a duplicate of the issue [#57492](https://github.com/tensorflow/tensorflow/issues/57492). Could you please check this issue, and confirm the same? Thank you!
 Hi, 
First of all, I must warn you that I am not an expert on the subject, so I apologize in advance if my answer lacks any coherence.

In my opinion the problems with your code are the following:

1. apparently the tf.cond() function executes both internal functions (true_fn and false_fn), somehow mapping the results and selecting them after checking the "pred" or flag.
2. flat_map() tries to return a dataset but one of the routes returns nothing (false_fn), and even worst, this function cuts the execution with an assert. Thus generating a conflict

The behavior of your code is then to fail when executing the false_fn branch and after that executes the true_fn since the predicate is a constant with value true.


Here my solutions:

By using tf.cond
``` python
import tensorflow as tf
import tensorflow_datasets as tfds

def do_not_come_here(image, flag):
    tf.print(flag)
    return image

def g(image, flag):
    return tf.cond(flag, lambda: image, lambda: do_not_come_here(image, flag))

def f(record):
    image = record['image']
    flag = tf.constant(True, dtype=tf.bool)
    processed_image = g(image, flag)
    processed_dataset = tf.data.Dataset.from_tensor_slices({
        'processed_image': processed_image
    })
    return processed_dataset

train, info = tfds.load("voc", split='train', with_info=True)
dataset = train.take(1).flat_map(f)

#To display the tf.print()
for example in dataset:
    pass
```

By using tf.py_function
``` python
import tensorflow as tf
import tensorflow_datasets as tfds

def do_not_come_here(image, flag):
    tf.print(flag)
    tf.debugging.assert_equal(0, 1)
    return image

def g(image, flag):
    return tf.py_function(
        func=lambda img, flg: img if flg else do_not_come_here(img, flg),
        inp=[image, flag],
        Tout=tf.uint8
    )

def f(record):
    image = record['image']
    flag = tf.constant(True, dtype=tf.bool)
    processed_image = g(image, flag)
    processed_dataset = tf.data.Dataset.from_tensor_slices({
        'processed_image': processed_image
    })
    return processed_dataset

train, info = tfds.load("voc", split='train', with_info=True)
dataset = train.take(1).flat_map(f)
``` Hi @tilakrayal 
> @eml-39502, Looks like this is a duplicate of the issue [#57492](https://github.com/tensorflow/tensorflow/issues/57492). Could you please check this issue, and confirm the same? Thank you!

I think the issue 57492 is having ```InvalidArgumentError: Condition x > y did not hold element-wise:``` exception. In my issue, I am casing the exception to confirm TF execution came to the place it should not have come to.

The issue here is ```tf.cond``` executes both ```true_fn``` and ```false_fn``` regardless with the predicate, hence it is not working as conditional branching function.

Hope this clarifies.
 @eml-39502,
Have you got the chance to have a look at the provided alternative workaround where the mentioned code was executed without any issue/error. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/c0261efbbd7cafd4ca48e703d03e7312/untitled1556.ipynb) and also the **tf.cond()** function evaluates both the `true_fn` and `false_fn`, regardless of the value of the pred argument. The resulting tensors are then mapped to the appropriate output tensor based on the value of pred. Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62438">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62438">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.14.1

### Custom code

No

### OS platform and distribution

NAME="Pop!_OS" VERSION="22.04 LTS" (Ubuntu equivalent)

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.8

### GPU model and memory

_No response_

### Current behavior?

The code below somehow always goes to ```do_not_come_here```. I suppose the usage ```tf.cond``` in the function ```g``` is not correct but not sure how to fix. Please advise.


## Code ##
```
import tensorflow as tf
import tensorflow_datasets as tfds


def do_not_come_here(image, flag):
    tf.print(flag)
    tf.assert_equal(0, 1)   

def g(image, flag):
    tf.cond(
        pred=flag,             # same with pred=tf.math.equal(flag, tf.constant(True))
        true_fn=lambda: image,
        false_fn=lambda: do_not_come_here(image, flag)
    )

def f(record):
    image = record['image']
    g(image, tf.constant(True))

train, info = tfds.load("voc", split='train', with_info=True)
train.take(1).flat_map(f)   # same with .map(f)
```

## Stacktrace

```
~/venv/ml/bin/python ~/reproduction.py 
2023-11-20 17:19:25.390562: I tensorflow/core/util/port.cc:111] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-11-20 17:19:25.415972: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-11-20 17:19:25.415993: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-11-20 17:19:25.416010: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-11-20 17:19:25.420375: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-11-20 17:19:26.669052: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.672374: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.672502: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.673349: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.673435: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.673493: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.729836: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.729955: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.730043: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.730112: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 1228 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4050 Laptop GPU, pci bus id: 0000:01:00.0, compute capability: 8.9
Traceback (most recent call last):
  File "~/reproduction.py", line 21, in <module>
    train.take(1).flat_map(f)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 2313, in flat_map
    return flat_map_op._flat_map(self, map_func, name=name)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/flat_map_op.py", line 24, in _flat_map
    return _FlatMapDataset(input_dataset, map_func, name)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/flat_map_op.py", line 33, in __init__
    self._map_func = structured_function.StructuredFunctionWrapper(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py", line 265, in __init__
    self._function = fn_factory()
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1222, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1192, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 694, in _initialize
    self._concrete_variable_creation_fn = tracing_compilation.trace_function(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 178, in trace_function
    concrete_function = _maybe_define_function(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 284, in _maybe_define_function
    concrete_function = _create_concrete_function(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 308, in _create_concrete_function
    traced_func_graph = func_graph_module.func_graph_from_py_func(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/framework/func_graph.py", line 1059, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 597, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py", line 231, in wrapped_fn
    ret = wrapper_helper(*args)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py", line 161, in wrapper_helper
    ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py", line 693, in wrapper
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.InvalidArgumentError: in user code:

    File "~/reproduction.py", line 18, in f  *
        g(image, tf.constant(True))
    File "~/reproduction.py", line 10, in g  *
        tf.cond(
    File "~/reproduction.py", line 7, in do_not_come_here  *
        tf.assert_equal(0, 1)

    InvalidArgumentError: Condition x == y did not hold element-wise:
    x (cond/assert_equal_1/x:0) = 
    0
    y (cond/assert_equal_1/y:0) = 
    1


Process finished with exit code 1
```



Same symptom with:
```
def g(image, flag):
    tf.cond(
        pred=tf.math.equal(flag, tf.constant(1)),
        true_fn=lambda: image,
        false_fn=lambda: do_not_come_here(image, flag)
    )

def f(record):
    image = record['image']
    g(image, tf.constant(1))
```

Same symptom with:
```
def g(image, flag):
    tf.cond(
        pred=tf.math.equal(flag, tf.constant(1)),       
        true_fn=lambda: do_not_come_here(image, flag), # <--- swap true_fn with false_fn
        false_fn=lambda: image
    )

def f(record):
    image = record['image']
    g(image, tf.constant(1))

```

## Environment

```
Python 3.10.12
tensorflow                   2.14.1
tensorflow-datasets          4.9.3
NAME="Pop!_OS" VERSION="22.04 LTS"
```

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import tensorflow_datasets as tfds


def do_not_come_here(image, flag):
    tf.print(flag)
    tf.assert_equal(0, 1)   

def g(image, flag):
    tf.cond(
        pred=flag,             # same with pred=tf.math.equal(flag, tf.constant(True))
        true_fn=lambda: image,
        false_fn=lambda: do_not_come_here(image, flag)
    )

def f(record):
    image = record['image']
    g(image, tf.constant(True))

train, info = tfds.load("voc", split='train', with_info=True)
train.take(1).flat_map(f)   # same with .map(f)
```


### Relevant log output

```shell
~/venv/ml/bin/python ~/reproduction.py 
2023-11-20 17:19:25.390562: I tensorflow/core/util/port.cc:111] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-11-20 17:19:25.415972: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-11-20 17:19:25.415993: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-11-20 17:19:25.416010: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-11-20 17:19:25.420375: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-11-20 17:19:26.669052: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.672374: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.672502: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.673349: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.673435: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.673493: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.729836: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.729955: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.730043: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2023-11-20 17:19:26.730112: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 1228 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4050 Laptop GPU, pci bus id: 0000:01:00.0, compute capability: 8.9
Traceback (most recent call last):
  File "~/reproduction.py", line 21, in <module>
    train.take(1).flat_map(f)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 2313, in flat_map
    return flat_map_op._flat_map(self, map_func, name=name)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/flat_map_op.py", line 24, in _flat_map
    return _FlatMapDataset(input_dataset, map_func, name)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/flat_map_op.py", line 33, in __init__
    self._map_func = structured_function.StructuredFunctionWrapper(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py", line 265, in __init__
    self._function = fn_factory()
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1222, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1192, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 694, in _initialize
    self._concrete_variable_creation_fn = tracing_compilation.trace_function(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 178, in trace_function
    concrete_function = _maybe_define_function(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 284, in _maybe_define_function
    concrete_function = _create_concrete_function(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 308, in _create_concrete_function
    traced_func_graph = func_graph_module.func_graph_from_py_func(
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/framework/func_graph.py", line 1059, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 597, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py", line 231, in wrapped_fn
    ret = wrapper_helper(*args)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py", line 161, in wrapper_helper
    ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
  File "~/venv/ml/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py", line 693, in wrapper
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.InvalidArgumentError: in user code:

    File "~/reproduction.py", line 18, in f  *
        g(image, tf.constant(True))
    File "~/reproduction.py", line 10, in g  *
        tf.cond(
    File "~/reproduction.py", line 7, in do_not_come_here  *
        tf.assert_equal(0, 1)

    InvalidArgumentError: Condition x == y did not hold element-wise:
    x (cond/assert_equal_1/x:0) = 
    0
    y (cond/assert_equal_1/y:0) = 
    1


Process finished with exit code 1
```
```
