tf.keras.ops.grad

Return a function that computes the gradient of f.

tf.keras.ops.grad(
    f, argnums=0
)

Experimental. This works on concrete tensors only, not symbolic ones, and raises NotImplementedError on backends without automatic differentiation. A non-scalar output is summed before differentiating, the same as a gradient tape would.

f A function returning a tensor. The arguments at argnums may be tensors, variables or nested structures of them.
argnums An integer or a tuple of integers, the positions of the arguments to differentiate with respect to. Defaults to 0.

A function with the signature of f that returns the gradient for each position in argnums, a single value for an integer and a tuple for a tuple. Arguments that do not affect the output get zeros.

Example:

def f(x, y):
    return x * y

x = ops.array([1.0, 2.0])
y = ops.array([3.0, 4.0])
dx = ops.grad(f)(x, y)
# dx is [3.0, 4.0]
dx, dy = ops.grad(f, argnums=(0, 1))(x, y)
# dy is [1.0, 2.0]
ops.grad(ops.tanh)(x)
# the elementwise derivative of tanh at x