length cannot be a NN parameter

[RTSandberg]

I have been unable to let length be an optimizable Pytorch.nn.Parameter without corrupting the transfer map. I think it is because every element inherits a length from the Element class.

cheetah/cheetah/accelerator/element.py

Line 32 in 3fe16a7

length: torch.Tensor = torch.zeros((1))

If I comment this out, then I am able to assign length as a Parameter, but this then causes some tests to fail (I believe it is every CI test using transfer maps).

I haven't found a good fix yet but it would be nice to allow length to be a parameter

[cr-xu]

Hi, indeed the length is causing quite some trouble. nn.Parameter doesn't like to re-register some attribute which is already defined in the class.

I implemented a quick and dirty way to circumvent this (see the new branch).
Now there should be no problem defining the length as a trainable nn.Parameter, and the tests are passing.

The fix was basically moving the default length definition into the __init__.
The issue with that is that Segment actually has a different length definition since it's adding all the lengths from its elements.
I now added _initialize_length in the base Element class and redefine it in Segment so that it doesn't clash with the length being a property.

@jank324 I'm not sure if it's the best solution though, any opinion?

[jank324]

Hmmm ... yeah ... we already struggled with this in #143 (which was part of #116).

Back then, some elements like Marker didn't have a length. Conceptually, I think this makes a lot of sense. But we had issues with broadcasting to the correct dimensions in situations where we needed to assume a length. I think an example of this would be transfer map reduction.

We decided as a fix just to give the Element class a length such that every element has one, that is Markers for example simply have length = 0.0. Conceptually, I think this is not as clean, but it still makes sense. This does not play nicely with the registration of nn.Parameter, which I believe is the only reason why we didn't allow length to be optimisable.

I think that now there are new considerations for this with #142. Here it makes even less sense for the element to have a length property. In fact, I think it would be too easily confused with what will probably be called something like effect_length.

A possible solution for this would be to just let every element define its own length again, i.e. no length in the base class. In my opinion this would be the cleanest solution. We would just have to find a good solution for #143. Admittedly, the original solution we found for this may have been just that little bit too hacky. 🙄

... so probably reverse and find a cleaner alternative for 9b9abe4.

[RTSandberg]

It sounds like you have thought through this some already and it's a bit of a thorny issue; I look forward to the fix!