Using TensorFlow autograph in NengoDL

Can I use autograph while building the TensorFlow computation graph implementing some NengoDL OpBuilder's build_step() function?

For example, if I want conditional execution of part of my code, am I limited to using tf.cond() to insert a decision node or can I use Python declarative instructions and convert using @tf.function?

For example, will this traditional TensorFlow code:

@Builder.register( MyLearningRule )
class MyLearningRuleBuilder( OpBuilder ):
    ...
    def build_step(self, signals):
        def if_error_over_threshold():
            # do something
        tf.cond( tf.reduce_any( tf.greater( tf.abs( local_error ), error_threshold ) ),
        true_fn=if_error_over_threshold,
        false_fn=lambda: None )

be equivalent in NengoDL to this one using autograph:

@Builder.register( MyLearningRule )
class MyLearningRuleBuilder( OpBuilder ):
    ...
    @tf.function
    def build_step(self, signals):
        if tf.reduce_any( tf.greater( tf.abs( local_error ), error_threshold ) ):
            # do something
        else:
            # do nothing

?

Hi Tioz,

Welcome to the Nengo forum.

Currently in NengoDL we do all the TensorFlow building in graph mode, not eager mode, so the autograph/tf.function functionality won’t work. We do plan on switching to eager mode soon, now that the performance issues have been largely resolved, at which point all that autograph stuff should work.

Excellent, I’m looking forwards to it!