Adds TensorCollection

[coreylowman]

Resolves #435
Related to #460

• Removes trait GradientUpdate
• Removes trait ParamUpdater
• optimizers now implement VisitTensors and use RecursiveWalker in their Optimizer impl
• Moves trait ResetParams to use TensorCollection, and removes all nn impls
• Adds NumParams using TensorCollection
• Refactor SaveToNpz/LoadFromNpz to use TensorCollection
• Removes all missing gradients tests in nn
  • Now that all nn layers use TensorCollection, testing npz/update/reset_params is enough to cover module walking
  • Unused tensor tests are now added to optimizers to ensure they properly capture unused
• Adds traits VisitsTensors
• Adds trait ModuleVisitor
• Adds trait TensorCollection

TODOs

• add documentation
• improve use statements
• Move visitors somewhere else other than tensor? top level?

[coreylowman]

@nkoppel thoughts? was able to impl optimizers with this approach as well

[nkoppel]

One thing I don't like about the current implementation is the repeated code in visitors/base.rs. A lot of this repetition could be removed by generalizing the behavior of accessing fields in &T, &mut T, and (&mut T, &T) with something like the following:

pub trait VisitTensors<E: Dtype, D: DeviceStorage> {
    type Argument<'a, M>: VisitTensorArgument<M> where M: 'a;
    type Err;

    fn visit<S: Shape>(
        &mut self,
        full_path: String,
        opts: TensorOptions<S, E, D>,
        t: Self::Argument<'_, Tensor<S, E, D>>,
    ) -> Result<(), Self::Err>;
}

pub trait VisitTensorArgument<M> {
    type WithModule<'a, Mod>: VisitTensorArgument<Mod> where Self: 'a;

    fn get_field<Field, GetRef, GetMut>(
        &mut self,
        get_ref: GetRef,
        get_mut: GetMut,
    ) -> Self::WithModule<'_, Field>
     where
        GetRef: FnMut(&M) -> &Field,
        GetMut: FnMut(&mut M) -> &mut Field,
        Field: TensorCollection<E, D>,
}

This would remove the need to have three different VisitTensor traits, but would require that we add extra constraints to the RecursiveWalker struct to ensure that it has F::Argument<M> as it's m field. Thoughts?

[coreylowman]

I totally agree, and nice sketch. Want to open a PR into this branch with that implemented? I think as long as we can keep that isolated to the implementation for RecursiveWalker (and not infect the user facing TensorCollection/TensorVisitor) it sounds like a great add

[coreylowman]

@nkoppel I'm pretty happy with where this is, any other organization/naming stuff you think we should do before merge?

[nkoppel]

@usbalbin merging this will break parts of #437. You'll need to remove the GradientUpdate implementation and replace it with a TensorCollection implementation.