Block masking in TensorFlow

Bhack · March 21, 2022, 1:18pm

This will already work also with jit_compile.

If you are interested we are tracking the nested conditional workaround in this example at:

github.com/tensorflow/tensorflow

tf.function condition with tensor

opened 10:49PM - 17 Mar 22 UTC

closed 05:00PM - 07 Apr 22 UTC

bhack

stat:awaiting tensorflower type:bug comp:autograph comp:tf.function TF 2.8

**System information** - Have I written custom code (as opposed to using a stoc…k example script provided in TensorFlow): Yes - OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab - Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: - TensorFlow installed from (source or binary): colab - TensorFlow version (use command below): default colab - Python version: default colab - Bazel version (if compiling from source): - GCC/Compiler version (if compiling from source): - CUDA/cuDNN version: - GPU model and memory: You can collect some of this information using our environment capture [script](https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh) You can also obtain the TensorFlow version with: 1. TF 1.0: `python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"` 2. TF 2.0: `python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"` **Describe the current behavior** ```python @tf.function def test(): b = tf.math.maximum(0,1) if 0 < b <= 2: b = b+1 return b test().numpy() ``` ```python OperatorNotAllowedInGraphError: in user code: File "<ipython-input-80-cfa5312e70c9>", line 4, in test * if 0 < b <= 2: OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature. ``` **Describe the expected behavior** I suppose that it will give `2` as with this little modified version: ```python @tf.function def test(): b = tf.math.maximum(0,1) if 0 < b and b <= 2: b = b+1 return b test().numpy() ``` ```python 2 ``` **[Contributing](https://www.tensorflow.org/community/contribute)** - Do you want to contribute a PR? (yes/no): With some hints - Briefly describe your candidate solution(if contributing): I need some pointer **Standalone code to reproduce the issue** Provide a reproducible test case that is the bare minimum necessary to generate the problem. If possible, please share a link to Colab/Jupyter/any notebook. **Other info / logs** Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

Wrapping the __call__ with autovectorization @tf.function with “exactly” the same code in the example it will require to cover additional builtins that we are tracking in the same ticket.
Or in the meantime you will need to substitute Python builtin min with an equivalent TF op.
See also:
https://tensorflow-prod.ospodiscourse.com/t/you-cannot-build-your-model-by-calling-build-if-your-layers-do-not-support-float-type-inputs/8356/6