Conv2D initialize error in tf function

[txir-met]

Hi, I am trying to write a gan with a convnet component with sonnet.

This is my base Conv function (there's some extra functionality to it but I'm trying to give a minimal example)

class Conv2D(snt.Conv2D):
      def __init__(self, *args, **kwargs):
        super().__init__(*args,**kwargs)
      def __call__(self, tensor):
        output = super().__call__(tensor)
        return(output)

This is the simplified step function:

@tf.function
def step(batch_inputs,batch_targets,generator_obj,gen_optimizer):
    return gen_step(batch_inputs,batch_targets,generator_obj,gen_optimizer)

and this is the simplified gen_step function:

def gen_step(batch_inputs, batch_targets, generator_obj,gen_optimizer):
  with tf.GradientTape() as tape:
    batch_predictions = generator_obj(batch_inputs)
  return batch_predictions

When I run this I get an error code from Conv2d. Here is the trace:

/lustre/home/mo-txirouch/google/layers.py:73 __call__  *
        output = super().__call__(tensor)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/utils.py:27 _decorate_unbound_method  *
        return decorator_fn(bound_method, self, args, kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/base.py:272 wrap_with_name_scope  *
        return method(*args, **kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/conv.py:117 __call__  *
        self._initialize(inputs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/utils.py:27 _decorate_unbound_method  *
        return decorator_fn(bound_method, self, args, kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/once.py:93 wrapper  *
        _check_no_output(wrapped(*args, **kwargs))
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/utils.py:27 _decorate_unbound_method  *
        return decorator_fn(bound_method, self, args, kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/base.py:272 wrap_with_name_scope  *
        return method(*args, **kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/conv.py:143 _initialize  *
        self.w = self._make_w()
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/utils.py:27 _decorate_unbound_method  *
        return decorator_fn(bound_method, self, args, kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/base.py:272 wrap_with_name_scope  *
        return method(*args, **kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/sonnet/src/conv.py:167 _make_w  *
        return tf.Variable(self.w_init(weight_shape, self._dtype), name="w")
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:262 __call__  **
        return cls._variable_v2_call(*args, **kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:256 _variable_v2_call
        shape=shape)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /home/mo-txirouch/.conda/envs/google/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:769 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

From prints I gathered the @tf.function runs again, presumably because of the autographing, which in turn causes the generator to rerun and leads to this error.

I don't know if this is an issue with the convolution in the sonnet library or with my implementation/workflow but after hacking at it for several days any advice would be appreciated. Thank you!