How to get symbolic learning_phase? (TF2 Eager)

[OverLordGoldDragon]

K.learning_phase() fetches the value, not the tensor itself - following backend.py, I found somewhat of a workaround, but it isn't user/API-friendly. I need the learning phase tensor to feed to K.function to get layer gradients, outputs, etc. Works fine w/ import keras.backend as K, but fails for import tensorflow.keras.backend as K.

Passing the symbolic learning_phase into K.function yields:

ValueError: Cannot create an execution function which is comprised of elements 
from multiple graphs.

---

Minimal applied example:

import tensorflow.keras.backend as K
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
import numpy as np

ipt = Input((16,))
out = Dense(16)(ipt)
model = Model(ipt, out)
model.compile('adam', 'mse')

x = np.random.randn(32, 16)
model.train_on_batch(x, x)

grads = model.optimizer.get_gradients(model.total_loss, model.layers[-1].output)
grads_fn = K.function(inputs=[model.inputs[0], model._feed_targets[0], K.learning_phase()], 
                      outputs=grads)

Full error trace:

File "<ipython-input-2-7f74922d7492>", line 3, in <module>
  outputs=grads)
File "D:\Anaconda\envs\tf2_env\lib\site-packages\tensorflow_core\python\keras\backend.py", line 3773, in function
  return EagerExecutionFunction(inputs, outputs, updates=updates, name=name)
File "D:\Anaconda\envs\tf2_env\lib\site-packages\tensorflow_core\python\keras\backend.py", line 3670, in __init__
  base_graph=source_graph)
File "D:\Anaconda\envs\tf2_env\lib\site-packages\tensorflow_core\python\eager\lift_to_graph.py", line 249, in lift_to_graph
  visited_ops = set([x.op for x in sources])
File "D:\Anaconda\envs\tf2_env\lib\site-packages\tensorflow_core\python\eager\lift_to_graph.py", line 249, in <listcomp>
  visited_ops = set([x.op for x in sources])

AttributeError: 'int' object has no attribute 'op'

---

Partial workaround:

import tensorflow.keras.backend as K
from tensorflow.python.eager import context
from tensorflow.python.ops import array_ops
import weakref
from tensorflow.python.framework import func_graph

def symbolic_learning_phase():
  graph = get_graph()
  with graph.as_default():
    if graph not in _GRAPH_LEARNING_PHASES:
      with K.name_scope(''):
        phase = array_ops.placeholder_with_default(
            False, shape=(), name='keras_learning_phase')
      _GRAPH_LEARNING_PHASES[graph] = phase
    return _GRAPH_LEARNING_PHASES[graph]

def get_graph():
  if context.executing_eagerly():
    global _GRAPH
    if _GRAPH is None:
      _GRAPH = func_graph.FuncGraph('keras_graph')
    return _GRAPH
  else:
    return ops.get_default_graph()

_GRAPH = None
_GRAPH_LEARNING_PHASES = weakref.WeakKeyDictionary()

symbolic_learning_phase()
# <tf.Tensor 'keras_learning_phase:0' shape=() dtype=bool>

[ravikyram]

I have tried on colab with TF version 2.0 ,2.1.0-dev20191124 and was able to reproduce the issue.Please, find the gist here.However i am not seeing any issue with TF version 1.15. Thanks!

[ymodak]

Duplicate #34508. Tracking this issue in another thread. Thanks!

[tensorflow-bot]

Are you satisfied with the resolution of your issue?
Yes
No