Example code output "Error: Gradient can only be implicitly created for scalar outputs"

[tomlarkworthy]

I tried the README code in a notebook: https://observablehq.com/d/e02d090ea2a71fda

but the grad() is throwing an error