Repository navigation
Remove generalized gramians #690
Copy link
Copy link
Closed
Labels
breaking-changeThis PR introduces a breaking change.This PR introduces a breaking change.cc: refactorConventional commit type for any refactoring, not user-facing, and not typing or perf improvementsConventional commit type for any refactoring, not user-facing, and not typing or perf improvementspackage: aggregationpackage: autogram
Description
The reasons are:
- This interface is not coherent with the interface of autojac. autojac.jac computes jacobians of shape [m, n_1, ..., n_l], even if the output if of shape [m_1, m_2], with m_1 * m_2 = m. So the coherent autogram behavior should be to compute gramians of shape [m, m] if the output is of shape [m_1, m_2].
- Changing autojac to make it coherent with autogram is extremely tedious: it would return jacobians of shape [m_1, ..., m_k, n_1, ..., n_l], which means that it's hard to know which dimensions are for the objectives and which dimensions are for the parameters. If we were to give such matrices to aggregators, we would also need to give them k, l or both. If we instead one day give jacobians of shape [m, n_1, ..., n_l] to aggregators, it will work without any problem.
- Even in autogram, this only works if the output is a single tensor. As we've discussed a while ago, if the output is e.g. two tensors of shape [m_1, m_2] and [m_3, m_4], finding the corresponding generalized gramians is a nightmare which we don't support and never will. Even if we managed to support that, making the corresponding generalized weightings would also be a nightmare, and explaining that to the users would basically be impossible.
- The concept of generalized gramian, GeneralizedWeighting and Flattening are extra classes that can confuse the user. If we also go down the path of having autojac coherent with autogram (i.e. by making it return jacobians of shape [m_1, ..., m_k, n_1, ..., n_l]) we would even need the concept of generalized jacobian, GeneralizedAggregators, and so on. It would become too much extremely quickly.
- It's actually possible to do IWMTL without all this. The user has the information of the shape of their output (e.g. [m_1, ..., m_k]), so if they really want to, they can themselves reshape their [m, m] gramian (again, with m = m_1 * ... * m_k) into a tensor of shape [m_1, ..., m_k, m_k, ..., m_1], and do whatever strategy they want to aggregate that (other than flattening of course, which they could have done without even doing this reshapping). So it's not like we're more expressive by supporting this: we're just taking a responsibility away from the user, but in a very incomplete way and while adding a lot of bloat to the library.
- In autogram's internals, we can still work with generalized gramians if it's easier, as long as we don't return them to the users.
So in my opinion, we should get rid of the concept of generalized gramian and PSDTensor (except maybe internally in autogram), of GeneralizedWeighting, and of Flattening, and make the necessary transition (update IWMTL example, tell users how to make the change in changelog, etc etc)
Metadata
Metadata
Assignees
Labels
breaking-changeThis PR introduces a breaking change.This PR introduces a breaking change.cc: refactorConventional commit type for any refactoring, not user-facing, and not typing or perf improvementsConventional commit type for any refactoring, not user-facing, and not typing or perf improvementspackage: aggregationpackage: autogram
Activity
PierreQuinton commentedon May 20, 2026
Not trivial, we should give an example of this and publish the associated Gramian reshape utility (which could return a Tensor an not a GeneralizedGramian).
Flattening is not useful, but I do not want to erase IWMTL from the library (at least not until we know it is a bad idea, which we don't, and which I believe might be the game changer we are waiting for).
Last thing, what would be the interface of autogram in case we have a batched dim and a non-batched dim? We input a matrix and it outputs a Gramian with [m1 m2, m1 m2] dimensions, that would correspond to the Gramian of the input matrix flattened? Not sure how to explain this in a non-confusing way.
ValerianRey commentedon May 20, 2026
I don't think it's our responsibility to help users with quite niche ideas of things they could do with their gramians. It's not like many people use autogram, let alone use it with outputs of more dimensions than 1, let alone wanna do something different than aggregating everything together (what Flattening was doing). If they do, they'll need custom code anyway. Helping them reshaping the gramian is only helping them with the first 20% of the job.
We don't need Flattening for IWMTL. Just compute the [m, m] gramian and aggregate it. So we'll keep the IWMTL example.
Yes. It's gonna be easier to explain than the current generalized gramians I'm sure. Just say that it's the gramian of the jacobian of the (flattened) outputs wrt the model params.