Use templates as parameters

[seanmor5]

Rather than requiring shape/type we can just use templates. The benefit is that this work supports using any Nx.Container as a model parameter. For right now we still support the old style under the param function, but this defers to the new parameter which requires templates.

There are certain places where using a container/composite parameter in place of a regular parameter makes sense. For example, if we want to initialize quantized models then we can create a quantized tensor container which represents a quantized parameter that can be converted back to a regular parameter.

In the future, I plan to unify the currently separated shape calculation and initialization process for a parameter. Realistically each parameter should just take input templates, do the shape calculation, and then initialize directly without any intermediate steps