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)