Add rewrite_nodes function

[seanmor5]

Resolves #529. This is a proper node rewriting API that should be much more flexible and easier than the wrap_nodes proposed earlier. Now LoRA can be implemented like this:

lora_dense_rewriter = fn [%Axon{} = x], %Axon{} = wx ->
  x
  |> Axon.dropout(rate: 0.1)
  |> Axon.dense(units_a, use_bias: false, name: "lora_a")
  |> Axon.dense(units_b, use_bias: false, name: "lora_b")
  |> Axon.multiply(scaling)
  |> Axon.add(wx)
end

lora_model = Axon.rewrite_nodes(model, fn
  %Axon.Node{op: :dense} -> lora_dense_rewriter
  _ -> :skip
end)

It also supports more complex rewriters. I will continue expanding this API until we have something closer to torch.fx