Comment décorer une fonction qui prend un tf.variable comme paramètre avec tf.function et surtout en utilisant la signature d'entrée

J'ai un problème où je dois modifier une variable dans une fonction Tensorflow. Ensuite, je dois convertir cette fonction en un graphique tensorflow. Le problème est que la taille de la variable n'est pas fixe. Exemple : il peut s'agir d'un ténosr de forme (3,) ou (2,). C'est pourquoi la fonction prend cette variable en paramètre, pour pouvoir la modifier et la retourner.

Voici un exemple de classe qui contient un appel de fonction , cette fonction prend deux arguments (x,v). x est un Tf.tenseur et v est un tf.Variable. v est affecté à la multiplication de x*v.

import tensorflow as tf

class MyModule(tf.Module):
  def __init__(self):
    pass

  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.int32), tf.TensorSpec(shape=[None], dtype=tf.int32)])
  def __call__(self, x, v):
    v.assign(x*v, read_value=False)
    return v

tf.config.run_functions_eagerly(False)
x = tf.constant([10,10])
v = tf.Variable(2*tf.ones_like(x), trainable=False)

module = MyModule()
module(x, v)

Cela fonctionne comme prévu en mode impatient, mais en mode graphique, j'obtiens l'erreur suivante : AttributeError : l'objet 'Tensor' n'a pas d'attribut 'assign'

Je sais que c'est à cause de la signature de tf.Variable. Ma question est comment puis-je spécifier la signature de tf.Variable étant donné que l'actuel produit une erreur?