How to use hook to clip a gradient?

General Usage
2

The error here is actually in your do block; it should be x, y, not (x, y), since as written it’s expecting a tuple of two things rather than two arguments.

However, that then returns (2, 4), which isn’t what you want. The reason is that hook applies to the gradient of the variable passed to it. This is slightly clearer if you write

gradient(2,3) do x, y
  l = m(x, y)
  Zygote.hook(clip, l)
end

clip gets passed l̄, the gradient of the loss, which is always 1, so this doesn’t do anything useful. To apply clip to (x, y) you have to pass that value to hook. Here’s one way to write that:

gradient(2,3) do x, y
  x, y = Zygote.hook(clip, (x, y))
  m(x, y)
end

If you can think of ways to make this stuff clearer in the docs a pr would be huge! This is definitely the kind of area we’d like to be a bit more comprehensive on.

1 Like