yes, otherwise the optimizer won't find the parameters to update. params in a torch.optim object are the IDs of the parameters.

This is a good point. As such, the implemented behaviour of not allowing to choose leave_parametrized=True whenever the size of the parametrized tensor is not the same as that of the original tensor is quite reasonable.

is this ever None? Isn't it initialized to an empty OrderedDict?

I think you're right. I'll delete this check.

return module for same behavior as before in pruning

This used to return a slightly more informative error when trying to remove pruning from an unpruned module ("Parameter '{}' of module {} has to be pruned before pruning can be removed" versus "No forward parametrizations found in tensor weight within module")

Changed the warnings thrown.

I would definitely split remove and undo into two separate functions given that they are two logically separate actions, as opposed to having: remove + leave pruned = remove reparametrization; remove + don't leave pruned = undo the pruning.

The one issue with undoing a reparametrization (and the reason why I chose not to implement an undo method for pruning) is that the definition of undoing is not clear for actions that can be called multiple times in an iterative fashion, like pruning. Will undo remove the whole history of pruning? or just the effect of the last pruning call? if the latter, then we would have to somehow keep track of all the history of mask changes, which instead I believe should be left to the user for specific use cases that require it. Another ambiguity around this is that once you make pruning permanent, you wouldn't be able to use undo anymore.

Happy to separate this one into two functions, I think that it can make it more clear from a user interface perspective.

I completely agree that, given the specific requirements of pruning, the undo method is tricky to implement on it. On the other hand, I think that the undo method can be quite useful in general. I can think of at least one general use case for this in the context of optimization on manifolds, where setting and undoing parametrizations has a very clear geometric meaning.

This has a side effect that now the weight gets serialized in the state dict, whereas before it wasn't. This may cause compatibility issues for loading saved models.

True. How should we go about this? In my opinion, it makes more sense to have weight as a buffer, and having it serialized with the whole model, given that tensor hooks (I don't know about module hooks) are not serialized.

The original (unpruned) weight and the mask are everything you need to recreate the pruned weight by simply multiplying the two together [true in pruning, not sure about other use cases]. So, in the pruning case, there would be no need for the weight to also be stored in the state_dict for saving -- it would just be redundant info. And this would be another copy of the whole net if all layers are at least partially pruned.

In general, though, wouldn't most users want to make the parametrization permanent (through .remove) before serializing?

I agree that it is somehow redundant, but it is not entirely It is not entirely redundant when saving the module, since the tensor hooks, at the moment, are not serialized. As such, if the user wants to save and resume later the training, they will have to set the parametrization again. In this case, the saved buffer can be used as a check to see that the saving / loading has been performed correctly. My reasoning behind putting them as a buffer was that this way it is more orthogonal with the way pytorch works. For example, they are moved to GPU when the user chooses to move the whole model to GPU. Otherwise this would have to be done by hand.

I do not think that one would want to make parametrizations permanent before serializing in general. In the general case, a parametrization is a constraint on certain parameter. For this reason, one may want to save a checkpoint of a model and resume the training later. To be able to do so, one needs the parametrization and the original tensors to resume the optimization with the constraint.

Doesn't a module already store its parametrizations as hooks without the need to pass them in separately as args?

Yep! Corrected now.

If users want to apply a parametrization that is logically compatible with pruning, they should be able to

This comment is outdated, I will change it. We do not support several forward parametrizations at the same time. I thought that these could be supported, but by the very nature of how forward hooks at a module level work, there is no obvious way to chain them. We allow to put several pruning parametrizations, because we are automatically merging them into the same PruningContainer.

This assumes a lot of interoperability between parametrizations. This is fine for pruning, because the pruning module handles how to exactly combine successive pruning calls, but this may not be true in the general reparametrization case.

It is left to the user to implement these in a compatible way. It does not seem wise to impose this artificially, given that there are plenty of use cases where this could be useful. For example, in the case of wanting to put constraints on a matrix, you may choose to have a constraint that makes the matrix positive definite, and then combine this with a constraint that makes the matrix positive definite with determinant one. This makes these very flexible and easy to mix and match.

needs an exception to raise :)

Why? It seems like a reasonable thing given that we allow registering the same parameter in multiple modules.

Many parametrizations are in general stateful (consider the mask in a pruning method, or the u and v tensors within spectral_norm), so it should not be allowed to register them on several tensors. This also allows for the current chain-like implementation (more on in a comment below)

Well but what if you really want to use the same parametrization from multiple modules? I don't see why should we forbid that. Parameters are stateful in their own ways too (e.g. .grad), but they still can be shared between modules.

I do not see any application for a shared parametrization. If a user really wanted to do so, they could create a new parametrization that shares the parameters with an existing one.

Having a parametrization only accepting a parameter allows writing parametrizations that can evaluate their current state calling self.evaluate(), and use this value to update their inner buffers, which is quite useful.

I don't think we should register the cache as a buffer. I'd like to avoid having those show up in state_dict() because they're not really meaningfully part of the module's state.

Fair enough.

weight_param seems like someone would be naming a "weight parameter". I think we should be a bit more verbose and call it a weight_parametrization.

I honestly don't love the mutable chaining syntax. My concerns are:

• It's generally better to avoid mutable state unless we really need it. It would be more natural to me to have it return a new parametrization that you can later register.
• The current syntax doesn't make it clear which parametrization gets evaluated first.
• How does it work if my parametrization has multiple parameters? Imagine a low-rank decomposition of the weight matrix as an example, as a cross product of two vectors.

Anyway, compare how mich clearer does this reads:

reagister_parametrization(module, "weight", p1 >> p2)

You can clearly see the order in which the parametrization pipeline is applied, and the operator can be made immutable. It doesn't solve the multiple parameter case though.

When you have multiple parameters, you can register them in a ModuleList and then implement the forward function as:

    def forward(self, Xs):
        return [p.evaluate(X) for p, X in zip(self.parametrizations, Xs)]

Funnily enough, I had already implemented a version of the lowrank example that you propose there, but with an SVD decomposition. In the case you propose, the parametrization should register an auxiliary parameter the size of the original vector, as the parametrization has more parameters than the original vector (it's a submersion).

The first and second points are very valid though. My reasoning to go with this design was the following: This chaining mechanism allows you, at any point, to evaluate the current value of the parametrization. I think of a parametrization as a RAII object over a buffer or a parameter, becoming the one that manages it. This way, if you have a chain of two parametrizations p1 from R^n to R^k and p2 from R^k to R^d over certain parameter t in R^n. In this setting, p2 does not have to know that it is acting on a parametrized object, it can always just ask for the value of the parameter, and it will be given p1(t), which lives in the correct space R^k, even though the original object doesn't.

This design also allows one to implement easily optimizations like the one currently being implemented in the pruning module, where if two pruning modules are put one after the other, their masks are merged into one by multiplying them. This could be done by simply overloading the chain method.

This is related to the point you made below of flattening this chain of parametrizations into something similar to a Sequential module. The problem with that is that each parametrization does not have a way to evaluate the current state of the tensor (which might be necessary to initialize the state of the parametrization, like the mask) and does not know anything about the previous parametrizations either (no optimization like that of pruning is possible).

When you have multiple parameters, you can register them in a ModuleList and then implement the forward function as:

Sorry, but I don't see how this applies to any of my points. Your example applies multiple independent parametrizations to multiple independent parameters, while I bring up issues about chaining parametrization on a single parameter.

Funnily enough, I had already implemented a version of the lowrank example that you propose there, but with an SVD decomposition. In the case you propose, the parametrization should register an auxiliary parameter the size of the original vector, as the parametrization has more parameters than the original vector (it's a submersion).

What does that mean? Where does it have to register this "auxiliary parameter"? I think all of the state necessary to construct a parametrization should be held entirely in the Parametrization object. It should never leak its own details into other modules, but that means that you cannot assume that there's a single underlying tensor you'll be processing.

The first and second points are very valid though. My reasoning to go with this design was the following: ...

While I see what you're saying, I don't understand how does it apply to my first two points. Chaining parametrizations doesn't require them to be stateful (just think of them as composition of stateless functions!), and the >> syntax can be overloaded too. The only difference is that you'd be implementing __rshift__ instead of chain in a subclass!

This design also allows one to implement easily optimizations like the one currently being implemented in the pruning module, ...

I don't see how my proposed design would prevent this. You can inspect the parametrizations in __rshift__ and can return different specialized composed forms if that can be done.

This is related to the point you made below of flattening this chain of parametrizations into something similar to a Sequential module. The problem with that is that each parametrization does not have a way to evaluate the current state of the tensor ... and does not know anything about the previous parametrizations either ....

Can you please elaborate on what do you mean when you say that "a parametrization doesn't have a way to evaluate the current state of the tensor"? Also, it's actaully good that it doesn't know anything about the next pamraetrizations, because that enables composability! The optimization you've mentioned should be applied whenever >> is applied. If the two parametrizations can be composed into a single one, then return the optimized version. If they can't, wrap them in the Sequential-like container. If you get a Sequential and another parametrization to append, then consider if you can't fuse the last two together and return a new Sequential with an optimized last layer.

On how to implement the low-rank layer: One of them would be to ask the low-rank parametrization to be applied on a tensor of size (2, n) and then implement

class LowRank(nn.Parametrization):
  def forward(self, t):
    a, b = t
    return a.unsqueeze(1) @ b.unsqueeze(0) # outer product

If the writer insists that the parametrization should be applied on a tensor, it is then when parameters in parametrizations come in to play. One would write:

class LowRankVector(nn.Parametrization):
    def __init__(self, n):
        super().__init__()
        self.b = nn.Parameter(torch.Tensor(1, n))

    def forward(self, a):
        return a.unsqueeze(1) @ self.b

Sometimes the latter option with parameters inside parametrizations is the only viable option if one wants to keep the same interface for different objects that implement different optimizations depending on the size of the input.

I do not think that we should support the possibility of applying parametrizations to two different parameters. If one wants to do so, they should couple them in just one parameter and decompose it inside forward. I have not found an example yet in which it is not reasonable to do so.

Given that parametrizations can be stateful, I do not think that returning a new parametrization is a very good option, as doing so would mean that you have to copy that parameter. In the case of pruning, for example, it would mean that one has to copy the mask and having a duplicate mask as well. I do not think that we want to deep-copy a module every time we do >>. To resolve the possible ambiguity of the order in which chain works we could rename it to compose. That way f.compose(g) would read and be semantically as f o g in mathematical notation.

Can you please elaborate on what do you mean when you say that "a parametrization doesn't have a way to evaluate the current state of the tensor

I completely agree that a parametrization should not know anything about the next parametrizations, but it should now about the previous ones. Now, consider the example of an adapter class, that takes a parametrization from R^n to R^n in the constructor and uses it to implement a parametrization form R^k to R^k via composing a function from R^k to R^n on the left and a function from R^n to R^k on the right. This case appears when trying to do orthogonal optimisation on rectangular matrices (R^{nxd}) based on optimization on squared orthogonal matrices (R^{nxn}). If one wants to do so, this adapter class has to prepend a parametrization to that going from R^k to R^n. When one has done so, the parametrization from R^n to R^n is ready to be used on elements of R^k. Furthermore, if it has functions that use the current value of the parametrization, it can just invoke self.evaluate() and, as it is aware of all the previous parametrizations that were chained to it, everything just works. If, on the other hand, we implement everything as a flat Sequential-like method, a parametrization that's in the middle could not evaluate all the parametrizations up to itself.

One of them would be to ask the low-rank parametrization to be applied on a tensor of size (2, n) and then implement

The assumption that you can pack all parameters of a parametrization into a single tensor is too strong. For starters, it's just not the right way to model this semantically. Secondly, what if you wanted to have a parametrization that would e.g. restrict a weight to be an affine transformation of another weight? You'd need to keep both the matrix and the bias, and doing that in a single tensor is super annoying.

I do not think that we should support the possibility of applying parametrizations to two different parameters.

I don't see why we should disallow that, especially that both parameters and modules support that. It would seem very incosistent to me if you coudn't do that.

Given that parametrizations can be stateful, I do not think that returning a new parametrization is a very good option, as doing so would mean that you have to copy that parameter.

I don't really follow your reasoning. The parametrization returned by >> can be a thin wrapper around existing parametrization objects, so there would be no need to reallocate them.

To resolve the possible ambiguity of the order in which chain works we could rename it to compose.

I don't think this resolves the issue. Both "chain" and "compose" have the same meaning and don't really clarify the direction in which the parametrizations are composed. Modelling it after a mathematical notation doesn't really help as I personally often struggle to recall which function comes first in f o g.

Furthermore, if it has functions that use the current value of the parametrization, it can just invoke self.evaluate() and, as it is aware of all the previous parametrizations that were chained to it, everything just works.

My question is why would it have functions that do that?

If, on the other hand, we implement everything as a flat Sequential-like method, a parametrization that's in the middle could not evaluate all the parametrizations up to itself.

I think that's a good thing, because it ensures compositionality!

Ok, I think there are two ways to look at composition of parametrizations. One is that it's simply function composition, and they don't have to be aware of each other. This naturally leads to a Sequential-like construction to represent their composition. The other way is to consider them more like a stack. But then, I'd argue that the right way to do it is to apply the parametrizations to parameters of existing parametrizations. This way you get the nice property that you can still call .evaluate() and get the desired result (the "previous" parametrizations will be automatically invoked when their respective attributes are accessed).

Does that make sense? I still don't fully understand why you might want to evaluate a parametrization and have it be aware of the things that come before it so it would be helpful to know. Unfortunately I wasn't able to follow your example. Also, if we go down the second path, then I'm pretty sure that you shouldn't be able to pass in custom tensors to .evaluate(). A parametrization should then simply have forward(self) (i.e. no arguments) that produces the result whenever called (and possibly caches the result).

Regarding your point of putting a parametrization that depends on several parameters, I have just implemented it so that one may write:
P.register_parametrization(module, ["weight", "bias"], MyParametrization(), name="combined")
where MyParametrization() is a parametrization whose forward takes two arguments. This way we allow to put parametrizations that take several inputs and return several outputs. Now the original tensors are not moved to the parametrization, but we just take a list of pointers to those tensors. The tensor will still be accessible under _tensor_name, that is, we prepend an underscore to its name.

I completely agree on the two ways of looking at parametrizations. Let me describe a more concrete, see if it makes sense. Assume that you have a function that goes from R^{n x n} to SO(n), where SO(n) is the space of the orthogonal matrices. This function reads in Pytorch:

def f(X):
  X = X.tril(-1)
  X = X - X.t()      # X is now skew-symmetric
  return expm(X) # The matrix exponential of a skew-symmetric matrix is orthogonal

Now, assume that we want to parametrize a non-square matrix of shape n x k to have orthogonal columns. It turns out that, if we append a matrix of size n x (n-k) of zeros on the left and pass it to f we get an n x n matrix. If we then take the k first columns of this orthogonal matrix, we get a matrix whose columns are orthogonal. This construction makes quite a bit of mathematical sense, but the why is not important. The idea is that there is an embedding that has to be applied before the f and a projection that has to be applied afterwards. This idea may be abstracted into a class that roughly reads:

class Fibration(P.Parametrization):
    def __init__(self, parametrization):
        self.parametrization = parametrization

    def embedding(self, X):
        raise NotImplementedError()

    def projection(self, X):
        raise NotImplementedError()

    def forward(self, X):
        X = self.embedding(X)
        X = self.parametrization.evaluate(X)
        X = self.projection(X)
        return X

Needless to say that represents a mathematical object of its own. Now, if we have a function inside parametrization that needs to make use of the evaluation of the current state of the parametrization chain, this implementation of Fibration will not work, as self.parametrization has not been registered anywhere.

On why does a parametrization need access to the parameters there are two examples. One is that of the initialization of parameters. In the RAII idea of a parametrization, they would take care of the initialization of their parameters, as they know how are they using them. For example, in the example of f, we are using tril to parametrize the skew-symmetric matrix, but we could also be using triu. As such, it is reasonable that the Parametrization that has f as its forward, is the one that exports methods to initialize the parameters that it manages. Here I have not explained why a parametrization would also need to evaluate the parametrization chain it is on, but I think that it follows reasonably from all this. I can continue this example to explain that if you still feel like it would help.

This RAII idea sort of enforces that a parametrization can only be registered in one parameter / tensor at a time. This idea is what the current implementation follows. If one wanted to effectively have the same parametrization on two parameters, one could create two different parametrizations and make them share their parameters.

Great, thanks for a more detailed write-up. To answer the first part (i.e. your example for why do we even need to be able to evaluate parametrizations on something else than its parameters), I think that we're mixing up two APIs that we really shouldn't. If you never need parameters, then you never should have made this thing a module in the first place. Ideally, you'd keep two APIs for parametrizations: one functional (where all parameters have to be passed in as arguments), and one nn.Module based, which takes care of state management for you, but generally defers to the stateless functional implementation in the forward. That's the pattern we already use for nn and we should continue it here.

On why does a parametrization need access to the parameters

I'm sorry, I can't follow your argument, because I'm not sure which part of my argumentation are you referring to. It would be very helpful if you could cite the relevant parts inline. Sorry 😕

By the way do you think it might be helpful to have a higher bandwidth meeting some time (like a voice call)? I feel like we've been going back and forth a lot, and writing up those arguments and rereading all the previous messages takes a lot of time. Can you please send me an email describing your availability so that we can arrange for a time? I'm hoping that would allow us to find a good design much more quickly!

Something to check before this gets merged: can parametrized models still get torch.saved? Pickle doesn't like weird things like this too much, so we should verify this.

I'm pretty sure it won't work out of the box, but it should be possible to fix it using some custom reducer functions.

I have never dealt with these corner cases of the pickling module, so I don't know how to go in this regard. Happy to take suggestion.

Can you please try pickling modules that have parametrizations applied and see if that fails or not? We'll think about solutions once we confirm that it is an issue.

Fails with:

_pickle.PicklingError: Can't pickle <class 'torch.nn.utils.parametrize.ParametrizedLinear'>: attribute lookup ParametrizedLinear on torch.nn.utils.parametrize failed

In the last commit I gave a go at this. Now parametrized layers work well with torch.load / torch.save. Take a look at it, see if you like it. I also added some testing for this.

This approach still throws a warning when being loaded:

UserWarning: Couldn't retrieve source code for container of type ParametrizedLinear. It won't be checked for correctness upon loading.
  warnings.warn("Couldn't retrieve source code for container of " 

I didn't manage to get rid of that one.

It is somewhat hacky, but one literally needs to hack around the limitations of pickle. If you know of a better way I can change it.

nit: orig is not a lot shorter than original and the latter reads much better. Please don't shorten the names too much.

I took the wording from the current implementation of the pruning module, but I agree. I'll change it.

Nevermind, I kept the name orig, because you also have the original() method that returns the original parameter from a parametrization chain. Should I change orig to _original?

Yeah let's do _original then.

This seems quite complicated, and I think it could be improved once you make Parametrization objects less stateful. You should be able to construct multiple of those, connect them into a container akin to nn.Sequential (maybe using this >> operator). Then you get this:

class ParametrizationChain(nn.Module):
  def __init__(self, *parametrizations):
    self.parametrizations = parametrizations

  def register(self, tensor):
    self.parametrizations[0].register(tensor)

  def evaluate(self, tensor=None):
    for p in parametrizations:
      tensor = p.evaluate(tensor)

Note that I've isolated the register function to make sure you don't have to much around too much with the orig attribute in register_parametrization. It should defer setting up of the state of the parametrization to the actual object.

Answered above.

I'm not so sure about the "append a prefix to the attribute" convention. Maybe it would be nicer to add something a'la ModuleDict called .parametrizations and then you'd be able to do linear.parametrizations.weight?

I took this convention from the current implementation of the pruning modules and spectral_norm and weight_norm. The problem with adding a new module dict is that one has to implement all the related stuff to iterate over it, and make it appear in __repr__ and so on, which is quite painful if this is implemented outside of Module. Maybe changing it to tensor_name + "_parametrization"?

Sorry can you elaborate on the module dict issues? What's the problem with iterating over it? Also, parameters usually don't appear in __repr__ so what's the deal there?

Parametrizations are modules, and might have parameters and buffers of their own. As such, they should appear in __repr__, and they should be returned by methods like children, modules, so that parameters works correctly.

I don't think that they necessarily have to appear in the __repr__ so I wouldn't bother with that. About appearing in all the submodule-related methods then yes, they should appear there, but you get that for free if you use ModuleDict, so what's the problem? It's a lot simpler than doing string arithmetic

Oh, I didn't know that those were iterated as sub-modules, but it does make a lot of sense. I just implemented it and works like a charm. It is also much cleaner.

I was thinking of doing the same and an OrderedDict for the caches. What do you think about that? Doing this we would completely get rid of all the nasty string arithmetic in the code, which I don't particularly like either.

I don't follow this caching logic. Why put it on the original module and not the Parametrization that has produced this? Also, if you put the caching logic into the Parametrization then the getter can always use the same code. Finally, it might be worth considering supporting setattr so that people can easily remove parametrizations by assigning a Parameter to a parametrized one, or overwriting the current one.