Using gen_nn_ops.max_pool_grad_v2 for layerwise relevance propagation not working: Keras symbolic input/output - tensorflow

Question

I am trying to apply layerwise relevance propagation (LRP) following some examples and tutorials, in particular this one and this one. However, I am struggling to get the relevances for the max_pooling layers. In particular, the problem appears when using gen_nn_ops.max_pool_grad_v2:
gen_nn_ops.max_pool_grad_v2(orig_input=a, orig_output=z, grad=s,
ksize=ksize,
strides=[1, 2, 2, 1], padding='VALID')
where the inputs are:
a : <KerasTensor: shape=(None, 6, 8, 512) dtype=float32 (created by layer 'activation_23')>
z : <KerasTensor: shape=(None, 3, 4, 512) dtype=float32 (created by layer 'tf.compat.v1.nn.max_pool_1')>
s: <KerasTensor: shape=(None, 3, 4, 512) dtype=float32 (created by layer 'tf.math.truediv_23')>
The error I got is:
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
/scratch/local/26230143/ipykernel_15814/2576584237.py in <module>
----> 1 gen_nn_ops.max_pool_grad_v2(orig_input=a, orig_output=z, grad=s,
2 ksize=kk,
3 strides=[1, 2, 2, 1], padding='VALID')
~/.local/lib/python3.9/site-packages/tensorflow/python/ops/gen_nn_ops.py in max_pool_grad_v2(orig_input, orig_output, grad, ksize, strides, padding, data_format, name)
6251 pass
6252 try:
-> 6253 return max_pool_grad_v2_eager_fallback(
6254 orig_input, orig_output, grad, ksize, strides, padding=padding,
6255 data_format=data_format, name=name, ctx=_ctx)
~/.local/lib/python3.9/site-packages/tensorflow/python/ops/gen_nn_ops.py in max_pool_grad_v2_eager_fallback(orig_input, orig_output, grad, ksize, strides, padding, data_format, name, ctx)
6283 data_format = "NHWC"
6284 data_format = _execute.make_str(data_format, "data_format")
-> 6285 _attr_T, _inputs_T = _execute.args_to_matching_eager([orig_input, orig_output, grad], ctx, [_dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.uint8, _dtypes.int16, _dtypes.int8, _dtypes.int64, _dtypes.bfloat16, _dtypes.uint16, _dtypes.half, _dtypes.uint32, _dtypes.uint64, ], _dtypes.float32)
6286 (orig_input, orig_output, grad) = _inputs_T
6287 ksize = _ops.convert_to_tensor(ksize, _dtypes.int32)
~/.local/lib/python3.9/site-packages/tensorflow/python/eager/execute.py in args_to_matching_eager(l, ctx, allowed_dtypes, default_dtype)
254 # not list allowed dtypes, in which case we should skip this.
255 if dtype is None and allowed_dtypes:
--> 256 tensor = ops.convert_to_tensor(t, ctx=ctx)
257 # If we did not match an allowed dtype, try again with the default
258 # dtype. This could be because we have an empty tensor and thus we
~/.local/lib/python3.9/site-packages/tensorflow/python/profiler/trace.py in wrapped(*args, **kwargs)
181 with Trace(trace_name, **trace_kwargs):
182 return func(*args, **kwargs)
--> 183 return func(*args, **kwargs)
184
185 return wrapped
~/.local/lib/python3.9/site-packages/tensorflow/python/framework/ops.py in convert_to_tensor(value, dtype, name, as_ref, preferred_dtype, dtype_hint, ctx, accepted_result_types)
1693
1694 if ret is None:
-> 1695 ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
1696
1697 if ret is NotImplemented:
~/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py in _constant_tensor_conversion_function(v, dtype, name, as_ref)
341 as_ref=False):
342 _ = as_ref
--> 343 return constant(v, dtype=dtype, name=name)
344
345
~/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py in constant(value, dtype, shape, name)
265 ValueError: if called on a symbolic tensor.
266 """
--> 267 return _constant_impl(value, dtype, shape, name, verify_shape=False,
268 allow_broadcast=True)
269
~/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py in _constant_impl(value, dtype, shape, name, verify_shape, allow_broadcast)
277 with trace.Trace("tf.constant"):
278 return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
--> 279 return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
280
281 g = ops.get_default_graph()
~/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py in _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
302 def _constant_eager_impl(ctx, value, dtype, shape, verify_shape):
303 """Creates a constant on the current device."""
--> 304 t = convert_to_eager_tensor(value, ctx, dtype)
305 if shape is None:
306 return t
~/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py in convert_to_eager_tensor(value, ctx, dtype)
100 dtype = dtypes.as_dtype(dtype).as_datatype_enum
101 ctx.ensure_initialized()
--> 102 return ops.EagerTensor(value, ctx.device_name, dtype)
103
104
~/.local/lib/python3.9/site-packages/keras/engine/keras_tensor.py in __array__(self, dtype)
252
253 def __array__(self, dtype=None):
--> 254 raise TypeError(
255 f'You are passing {self}, an intermediate Keras symbolic input/output, '
256 'to a TF API that does not allow registering custom dispatchers, such '
TypeError: You are passing KerasTensor(type_spec=TensorSpec(shape=(None, 6, 8, 512),
dtype=tf.float32, name=None), name='activation_23/Relu:0', description="created by layer
'activation_23'"), an intermediate Keras symbolic input/output, to a TF API that does not allow
registering custom dispatchers, such as `tf.cond`, `tf.function`, gradient tapes, or
`tf.map_fn`. Keras Functional model construction only supports TF API calls that *do* support
dispatching, such as `tf.math.add` or `tf.reshape`. Other APIs cannot be called directly on
symbolic Kerasinputs/outputs. You can work around this limitation by putting the operation in a
custom Keras layer `call` and calling that layer on this symbolic input/output.
I am assuming that somehow a, z and s must be tf.Tensor (am I wrong)? (instead of KerasTensor) , but I don't really know and neither how to get around this issue. I read that using:
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()
might work, but it is not working for me. I'd really appreaciate any idea, suggestion here!
Thank you very much in advance,