0

I'm implementing a custom layer in tensorflow 2.x . My requirement is such that, the program should check a condition before returning the output.

class SimpleRNN_cell(tf.keras.layers.Layer):
    def __init__(self, M1, M2, fi=tf.nn.tanh, disp_name=True):
        super(SimpleRNN_cell, self).__init__()
        pass        
    def call(self, X, hidden_state, return_state=True):
        y = tf.constant(5)
        if return_state == True:
            return y, self.h
        else:
            return y

My question is: should I continue using the present code (assuming that the tape.gradient(Loss, self.trainable_weights) will work fine) or should I use tf.cond(). Also, if possible please explain where to use tf.cond() and where not to. I haven't found much content on this topic.

desertnaut
  • 57,590
  • 26
  • 140
  • 166
Lawhatre
  • 1,302
  • 2
  • 10
  • 28

1 Answers1

2

tf.cond is only relevant when performing conditional evaluation based on data in the differentiable computation graph. (https://www.tensorflow.org/api_docs/python/tf/cond) This was especially necessary in TF 1.0 with graph mode being the default. For eager mode the GradientTape system allows to also do conditional data flow with the python constructs such as if ...: (https://www.tensorflow.org/guide/autodiff#control_flow)

Yet for just providing different behavior based on configuration parameters, that don't depend on data from the computational graph and are fixed during the model runtime, using simple python if statements is correct.

sholderbach
  • 96
  • 1
  • 5
  • Got it :) But what if ```y``` was not fixed in the code and was dependent on some calculation? so I don't need ```tf.cond``` unless I'm using ```@tf.function```. Is that what you mean? – Lawhatre Feb 14 '21 at 07:10
  • 1
    In principle most of that is also covered by the autograph features. https://www.tensorflow.org/guide/intro_to_graphs#flow_control_and_side_effects https://www.tensorflow.org/guide/function#conditionals But some restrictions apply: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#if-statements And side-effects could cause issues https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#effects-of-the-tracing-process – sholderbach Feb 14 '21 at 17:22