Initial draft of hooks

[seanmor5]

Resolves #77 when merged.

This is an initial draft but it motivates 2 discussions: host callbacks and reusable model blocks.

In order to perform side-effecting operations with hooks, Nx needs to support host callbacks. Once that's done, it likely means that hooks will always be implemented as host callbacks, but the implementation will need to evolve once we support them. I don't see the API being impacted much once we introduce host callbacks.

As for model blocks: In PyTorch has two kinds of hooks - tensor hooks and module hooks. In Axon we're more focused on module-style hooks for now. Module hooks have the following signature: hook(module, input, output) :: None. The equivalent in Axon would be: hook(axon_struct, input, output) :: nil. For both forward and backward hooks, even if they are only meant to be side-effecting operations, I'd prefer they always return an output such that after a layer the output is basically what's shown in this PR.

One issue is that we're currently unable to support hooks registered on large blocks of a model. For example, if I have a block of Convs:

block =
  x
  |> conv()
  |> conv()
  ...

and I tried to register a hook, I would only have access to the input to the last conv in the block, not x. This motivates introducing a module-like abstraction which allows you to treat groups of layers as a single block. This is essentially the same issue we have in #59 as well.

[josevalim]

Looks great! Just one note: I think even with host callbacks, you would want to choose if you want to bring the data back or the hook is a Nx transformation.

[seanmor5]

I will have a new PR for this shortly integrated with Nx's new hook API