Fix "AttributeError" thrown when processing multiple subgraphs induced by custom gradient definitions

[nothingstopsme]

This tiny patch is intended to fix the problem described in the title, which is probably best explained with an example as follows.

Considering this simple case:

import tensorflow as tf
import tf2jax

@tf.custom_gradient
def my_op(x):
  y = x**2
  def grad(dy):
    return dy * x * 2

  return y, grad

@tf.function
def func(z):
  return my_op(my_op(z))

with tf2jax.override_config("convert_custom_gradient", True):
  func_jaxed = tf2jax.convert_functional(func, np.zeros((2, 2), np.float32))

The conversion starts with generating a full graph depicted by func() and all op nodes involved:

nodes = [
  _OpNode(name='z'), _OpNode(name='pow/y'), _OpNode(name='pow'), _OpNode(name='IdentityN'),
  _OpNode(name='pow_1/y'), _OpNode(name='pow_1'), _OpNode(name='IdentityN_1'), _OpNode(name='Identity_2')
]

Since there are 2 calls to my_op(), which defines its own custom gradient, 2 subgraphs grouping associated nodes are extracted from the full graph:

# nodes belonging to _Subgraph(name='IdentityN')
[_OpNode(name='pow/y'), _OpNode(name='pow'), _OpNode(name='IdentityN')]

# nodes belonging to _Subgraph(name='IdentityN_1')
[_OpNode(name='pow_1/y'), _OpNode(name='pow_1'), _OpNode(name='IdentityN_1')]

Then the process goes on to update nodes to reflect the subgraph structures, supposedly resulting in:

nodes = [_OpNode(name='z'), _Subgraph(name='IdentityN'), _Subgraph(name='IdentityN_1'), _OpNode(name='Identity_2')]

Surprisingly, it turns out that tf2jax fails to convert this example and raises AttributeError at the converting stage above. This is because _Subgraph.rewrite() (defined in "tf2jax/_src/tf2jax.py"), which is the function responsible for manipulating nodes and subgraph substitution, treats every item in nodes as just a op node of the type _OpNode; however, this is not true after the addition of the 1st subgraph _Subgraph(name='IdentityN') to nodes, when the content of nodes becomes

# after "_Subgraph(name='IdentityN')" is substituted
nodes =  [
  _OpNode(name='z'), _Subgraph(name='IdentityN'), _OpNode(name='pow_1/y'),
  _OpNode(name='pow_1'), _OpNode(name='IdentityN_1'), _OpNode(name='Identity_2')
]

Therefore, as the subsequent call to _Subgraph.rewrite() on the 2nd subgraph _Subgraph(name='IdentityN_1') attempts to modify an attribute that is editable on _OpNode but read-only on _Subgraph while iterating nodes, an AttributeError exception occurs.

This fix takes into account the fact that nodes could be a mix of op nodes and subgraphs, and adds checks to avoid invalid attribute/property write attempts.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.