<p dir="auto"><strong>System information</strong></p>
<ul dir="auto">
<li>Have I written custom code (as opposed to using a stock example script provided in TensorFlow): no</li>
<li>OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Archlinux</li>
<li>TensorFlow installed from (source or binary): binary</li>
<li>TensorFlow version (use command below): v1.12.1-3374-g9eb67b17bf 2.0.0-dev20190605</li>
<li>Python version: 3.6</li>
<li>CUDA/cuDNN version: 10</li>
<li>GPU model and memory: 1080 Ti</li>
</ul>
<p dir="auto"><strong>Describe the current behavior</strong></p>
<p dir="auto">I expect to do a forward pass with a model with a BachNormalization layer in training mode, when using the <code class="notranslate">tf.distribuite.MirroredStrategy</code> but I can't, because it reises the following exception:</p>
<blockquote>
<p dir="auto">RuntimeError: <code class="notranslate">add_update</code> was called in a cross-replica context. This is not expected. If you require this feature, please file an issue.</p>
</blockquote>
<p dir="auto">Why it is not expected?</p>
<p dir="auto"><strong>Describe the expected behavior</strong></p>
<p dir="auto">It should work.<br>
The commit that introduced this behavior is: <a class="commit-link" href="https://github.com/tensorflow/tensorflow/commit/316cd57883166e6a0b4c2d0eaacebddad7675b39#diff-8eb7e20502209f082d0cb15119a50413"><tt>316cd57</tt>#diff-8eb7e20502209f082d0cb15119a50413</a></p>
<p dir="auto"><strong>Code to reproduce the issue</strong></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow as tf

model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(10),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dense(1),
    ]
)

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    out = model(tf.zeros((1, 10)), training=True)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>

<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-v">Sequential</span>(
    [
        <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-c1">10</span>),
        <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">BatchNormalization</span>(),
        <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-c1">1</span>),
    ]
)

<span class="pl-s1">strategy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">distribute</span>.<span class="pl-v">MirroredStrategy</span>()
<span class="pl-k">with</span> <span class="pl-s1">strategy</span>.<span class="pl-en">scope</span>():
    <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-en">model</span>(<span class="pl-s1">tf</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">1</span>, <span class="pl-c1">10</span>)), <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)</pre></div>
<h3 dir="auto">System information</h3>
<ul dir="auto">
<li><strong>Have I written custom code (as opposed to using a stock example script provided in TensorFlow)</strong>: No</li>
<li><strong>OS Platform and Distribution (e.g., Linux Ubuntu 16.04)</strong>:</li>
</ul>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="$ uname -mrs
Linux 4.14.0-49.el7a.ppc64le ppc64le"><pre class="notranslate">$ uname -mrs
Linux 4.14.0-49.el7a.ppc64le ppc64le</pre></div>
<ul dir="auto">
<li><strong>TensorFlow installed from (source or binary)</strong>: Binary</li>
<li><strong>TensorFlow version (use command below)</strong>: 1.6.0</li>
<li><strong>Python version</strong>: 2.7.5</li>
<li><strong>Bazel version (if compiling from source)</strong>:</li>
<li><strong>GCC/Compiler version (if compiling from source)</strong>:</li>
<li><strong>CUDA/cuDNN version</strong>: CUDA 9.1, CuDNN 7.0.5</li>
<li><strong>GPU model and memory</strong>: Tesla V100:</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="== nvidia-smi ===================================================
Thu Apr 26 18:35:34 2018       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 396.19                 Driver Version: 396.19                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000004:04:00.0 Off |                    0 |
| N/A   29C    P0    52W / 300W |      0MiB / 15360MiB |      0%   E. Process |
+-------------------------------+----------------------+----------------------+"><pre class="notranslate"><code class="notranslate">== nvidia-smi ===================================================
Thu Apr 26 18:35:34 2018       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 396.19                 Driver Version: 396.19                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000004:04:00.0 Off |                    0 |
| N/A   29C    P0    52W / 300W |      0MiB / 15360MiB |      0%   E. Process |
+-------------------------------+----------------------+----------------------+
</code></pre></div>
<ul dir="auto">
<li><strong>Exact command to reproduce</strong>: N/A</li>
</ul>
<p dir="auto">We would like to be able to have access to the CuDNN autotuner in <code class="notranslate">tf.keras</code> module to access optimal algorithms for a given hardware (or, perhaps passing a custom convolutional algorithm from config). In TensorFlow, I can specify to use the CuDNN autotuner by setting:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="os.environ['TF_CUDNN_USE_AUTOTUNE'] = &quot;1&quot;"><pre class="notranslate">os.environ[<span class="pl-s"><span class="pl-pds">'</span>TF_CUDNN_USE_AUTOTUNE<span class="pl-pds">'</span></span>] = <span class="pl-s"><span class="pl-pds">"</span>1<span class="pl-pds">"</span></span></pre></div>
<p dir="auto">(currently enabled by default), which improves performance significantly on Volta GPUs and especially with FP16.</p>
<p dir="auto">However, I am unable to access this performance improvement when running pure <code class="notranslate">tf.keras</code>, where setting this environmental variable does not have any effect.</p>
<h3 dir="auto">Source code / logs</h3>
<p dir="auto">Following simple example could be used to reproduce the issue: <a href="https://gist.github.com/ASvyatkovskiy/8d1dd622e447d9d8de1ec4e238e0dbaa">https://gist.github.com/ASvyatkovskiy/8d1dd622e447d9d8de1ec4e238e0dbaa</a></p>