title: Conditional branches using tf.function
body: <p>I am having problems calculating the gradient using the gradient tape when using a <code>tf.function</code> with conditional branches.</p>
  
  <p>Inside a gradient tape scope, I am trying to calculate the gradient of <code>z</code> w.r.t <code>self.LMn</code>. This works perfectly fine when I do not annotate the function with <code>@tf.function</code>. The error originates in a subclassed <code>tf.keras.layers.Layer</code> call function:</p>
  
  <pre><code>def call(self, x, training=None, labels=None, cur_switch=None, basis_filters=None):
      if basis_filters is None:
          basis_filters = self.initial_basis_filters
  
      gap = tf.reduce_mean(x, axis=[1, 2])
      z = tf.einsum('bc,cd-&gt;bd', gap, self.LMn)
  
      # ... code for out
  
      return out, z
  </code></pre>
  
  <p>The actual error is given as follows:</p>
  
  <pre><code>....
  C:\...\ops\cond_v2.py:387 &lt;lambda&gt;
      lambda: _grad_fn(func_graph, grads), [], {},
  C:\...\ops\cond_v2.py:363 _grad_fn
      assert len(func_graph.outputs) == len(grads)
  
  AssertionError: 
  </code></pre>
  
  <p>More specifically,</p>
  
  <pre><code>func_graphs.outputs = [&lt;tf.Tensor 'vgg16/block1a/block1a_conv/cond_2/Identity:0' shape=(128, 32, 32, None) dtype=float32&gt;, &lt;tf.Tensor 'vgg16/block1a/block1a_conv/cond_2/OptionalFromValue:0' shape=() dtype=variant&gt;, &lt;tf.Tensor 'vgg16/block1a/block1a_conv/cond_2/OptionalFromValue_1:0' shape=() dtype=variant&gt;]
  </code></pre>
  
  <p>and </p>
  
  <pre><code>grads = (&lt;tf.Tensor 'gradient_tape/vgg16/block1a/block1a_conv/strided_slice_3/StridedSliceGrad_1:0' shape=(128, 32, 32, None) dtype=float32&gt;,)
  </code></pre>
  
  <p>I can assume that each of these outputs corresponds to the cases for some set of conditional branches outputs that are then fed into the <code>tf.einsum</code> function. I have read over all of the edge cases and precautions in the gradient tape documentation as this seems to the problem. Just a note that I am only performing conditional computation using hyperparameters (passed into <code>tf.function</code> as pythonic variables, such as <code>basis_filters</code>). There are also some conditional branches using (tensorflow ops) functions of these hyperparameters, is this allowed? or do I need to compute these values outside <code>tf.function</code> and pass these in as pythonic variables too?</p>
  
  <p>I know the question is not completely clear and I can provide any extra information if needed. It would very helpful to have some guidance on what to look for with this kind of problem.</p>
  
  <p>Thanks!</p>