def model(a, b):
     with tape_lib.GradientTape() as tape:
         tape.watch(a)
         tape.watch(b)
         result = unified_math_ops.mul(a, b)
     grads = tape.gradient(result, [a, b])
     return grads
 def model(a, b):
     return unified_math_ops.mul(a, b)