feat(autogram): Add DiagonalSparseTensor.
#466
Draft
+1,245
−13
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
…ctions (mean and sum).
ValerianRey
force-pushed
the
block-diagonal-tensor
branch
from
87b4da0 to
efa8019
Compare
…hould be implemented differently)
…used over the constructor of `DST`
|
I think square is failing because it is implemented as |
…ake `data` and `v_to_p` public.
…e, and make `data` and `v_to_p` public." This reverts commit 85c8e41.
4 tasks
ValerianRey
reviewed
Oct 30, 2025
…k-diagonal-tensor
…ungrouped_dims, remove print
…irtual dimension that uses a physical dimension multiple times.
* The result is the same as before * Before that we only iterated on the pdims used by each virtual dim, and summed them if a pdim was present multiple times. * Now the new stride is already the sum of the old strides when a pdim is present multiple times in a vdim. We iterate over all dimensions, because for dimensions not present in the vdim, the stride is simply 0. * There's probably a more efficient implementation
* Now that we iterate over all_pdims instead of the pdims of the current virtual dimension, the result of torch.stack([p_indices_grid[d] for d in all_pdims], dim=-1) is always the same, and is simply equal to torch.stack(p_indices_grid, dim=-1). So we directly stack the p_indices_grid when creating it, and use the already stacked p_indices_grid in the for-loop.
* In the long term I'll try to rely mostly or even only on them instead of v_to_ps, so it makes sense to pre-compute them.
… make it directly as a tuple
* This makes lines shorter
…cannot have strides on cuda because addmm_cuda (required for tensordot) does not support Long tensors (but the cpu version does).
PierreQuinton
commented
Oct 31, 2025
| self.physical = physical | ||
| self.v_to_ps = v_to_ps | ||
| pshape = list(self.physical.shape) | ||
| self.strides = tensor([strides_v2(pdims, pshape) for pdims in self.v_to_ps]) |
There was a problem hiding this comment.
Could we make strides_v2 return a Tensor directly so that if we change from v_to_ps to pure stride description, we can then already assume it is an array (we could have a function to build a stride from a v_to_ps and physical.shape)
PierreQuinton
commented
Oct 31, 2025
| # addmm_cuda not implemented for Long tensors => gotta have these tensors on cpu | ||
| v_indices_grid = tensordot(self.strides, p_indices_grid, dims=1) | ||
| res = zeros(self.shape, device=self.device, dtype=self.dtype) | ||
| res[tuple(v_indices_grid)] = self.physical |
There was a problem hiding this comment.
I would be really surprised if the cast to tuple was necessary (but did not check)
Suggested change
| res[tuple(v_indices_grid)] = self.physical | |
| res[v_indices_grid] = self.physical |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Some prototype of diagonal sparse Tensors.
TODO:
contiguous_data/v_to_psmatch up to jointmove_dim.Problems and questions:
t=DST(..., v_to_p=[[0],[0]])and slicingt[2:8, 2:8], we would want the result to be aDST(..., v_to_p=[[0], [0]])but this cannot be done easily with slicing iteratively on each dimensions.atenon innerTensorrather than higher level operators fromtorch? It feels likeatendo not accept any Tensor and we would want to be able to compose with any. However doing this might result in some extra overhead.