Tensorflow Autograph issue with for loop

Question

I would like to extract the derivatives of a Tensorflow model with respect to its input using autograph for speed, but for some reason autograph gives an error. Here is the function that I am using:

@tf.function
def get_dydt(data,model,outsize):
    ipt = data
    y_list = []
    with tf.GradientTape(persistent=True,watch_accessed_variables=True) as tape_1:
        tape_1.watch(ipt)

        with tf.GradientTape(persistent=True, watch_accessed_variables=True) as tape_2:
            tape_2.watch(ipt)
            y = model(ipt)


        dydt_list = tf.TensorArray(tf.float32, size=0, dynamic_size=True,infer_shape = True)

        for i in tf.range(outsize):
            
            dydt_list.write(i,tape_2.gradient(y[i,:],ipt))
        dydt = dydt_list.stack()

    return dydt

This gives the following error: ValueError: None values not supported. for the line dydt_list.write(i,tape_2.gradient(y[:,i],ipt)). I assume this error is due to the fact that placeholders are used for the sizes, but I'm not sure how to do this in a way that autograph will accept. What can I do to address this?

Full traceback:

----------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-50-b524f1fd6255> in <module>
----> 1 training_step_test_2(batch,model)

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\def_function.py in __call__(self, *args, **kwds)
    826     tracing_count = self.experimental_get_tracing_count()
    827     with trace.Trace(self._name) as tm:
--> 828       result = self._call(*args, **kwds)
    829       compiler = "xla" if self._experimental_compile else "nonXla"
    830       new_tracing_count = self.experimental_get_tracing_count()

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\def_function.py in _call(self, *args, **kwds)
    869       # This is the first call of __call__, so we have to initialize.
    870       initializers = []
--> 871       self._initialize(args, kwds, add_initializers_to=initializers)
    872     finally:
    873       # At this point we know that the initialization is complete (or less

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\def_function.py in _initialize(self, args, kwds, add_initializers_to)
    723     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    724     self._concrete_stateful_fn = (
--> 725         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    726             *args, **kwds))
    727 

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   2967       args, kwargs = None, None
   2968     with self._lock:
-> 2969       graph_function, _ = self._maybe_define_function(args, kwargs)
   2970     return graph_function
   2971 

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\function.py in _maybe_define_function(self, args, kwargs)
   3359 
   3360           self._function_cache.missed.add(call_context_key)
-> 3361           graph_function = self._create_graph_function(args, kwargs)
   3362           self._function_cache.primary[cache_key] = graph_function
   3363 

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3194     arg_names = base_arg_names + missing_arg_names
   3195     graph_function = ConcreteFunction(
-> 3196         func_graph_module.func_graph_from_py_func(
   3197             self._name,
   3198             self._python_function,

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    988         _, original_func = tf_decorator.unwrap(python_func)
    989 
--> 990       func_outputs = python_func(*func_args, **func_kwargs)
    991 
    992       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\eager\def_function.py in wrapped_fn(*args, **kwds)
    632             xla_context.Exit()
    633         else:
--> 634           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    635         return out
    636 

~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\func_graph.py in wrapper(*args, **kwargs)
    975           except Exception as e:  # pylint:disable=broad-except
    976             if hasattr(e, "ag_error_metadata"):
--> 977               raise e.ag_error_metadata.to_exception(e)
    978             else:
    979               raise

ValueError: in user code:

    <ipython-input-49-24030d8d0c52>:55 training_step_test_2  *
        dydt_list.write(i,tape_2.gradient(y[i,:],ipt))
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\util\tf_should_use.py:247 wrapped  **
        return _add_should_use_warning(fn(*args, **kwargs),
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\ops\tensor_array_ops.py:1159 write
        return self._implementation.write(index, value, name=name)
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\ops\tensor_array_ops.py:534 write
        value = ops.convert_to_tensor(
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\profiler\trace.py:163 wrapped
        return func(*args, **kwargs)
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\ops.py:1540 convert_to_tensor
        ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\constant_op.py:339 _constant_tensor_conversion_function
        return constant(v, dtype=dtype, name=name)
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\constant_op.py:264 constant
        return _constant_impl(value, dtype, shape, name, verify_shape=False,
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\constant_op.py:281 _constant_impl
        tensor_util.make_tensor_proto(
    C:\Users\~\anaconda3\envs\tensorflowGPU\lib\site-packages\tensorflow\python\framework\tensor_util.py:445 make_tensor_proto
        raise ValueError("None values not supported.")

    ValueError: None values not supported.

This problem can be reproduced in the following manner:

net = tf.keras.models.Sequential()
net.add(tf.keras.Input(shape = (1,)))
for i in range(5):
    net.add(tf.keras.layers.Dense(50,activation = 'swish'))
net.add(tf.keras.layers.Dense(3))

ipt = tf.random.uniform([4,1],0,1)

get_dydt(ipt,net,3)

Answer 1

The error comes from the fact that you are trying to write a value of None in a tf.TensorArray. You get this error because your gradient calculation outputs None because some operations that you are doing are not done inside the scope of the tf.GradientTape. You can read that answer for more details: Answer to: "tensorflow differentiate only element of vector".

You need to slice your tensor in the scope of the GradientTape to get the actual gradient. One way to achieve this is the following:

@tf.function
def get_dydt(data,model,outsize):
    ipt = data
    dydt_list = tf.TensorArray(tf.float32, size=0, dynamic_size=True,infer_shape = True)

    with tf.GradientTape(persistent=True,watch_accessed_variables=True) as tape_1:
        tape_1.watch(ipt)

        with tf.GradientTape(persistent=True, watch_accessed_variables=True) as tape_2:
            tape_2.watch(ipt)
            y = model(ipt)
            # slicing the tensor in the scope of the tape
            y_slices = [y[i,:] for i in range(outsize)]

        for i, y_slice in enumerate(y_slices):
            grad = tape_2.gradient(y_slice,ipt)
            # you should reuse the returned Array by write to ensure 
            # that the writes occur.
            dydt_list = dydt_list.write(i,grad)
        dydt = dydt_list.stack()

    return dydt