Function f

unstable_source/vmap_recipe.py:85–86

(x)

# Setup
N = 5
def f(x):
    return x ** 2

x = torch.randn(N, requires_grad=True)
y = f(x)