Skip to content
Draft
Changes from 1 commit
Commits
Show all changes
146 commits
Select commit Hold shift + click to select a range
32ef75e
Add `DiagonalSparseTensor` with the default fallback to dense mechanism.
PierreQuinton Oct 20, 2025
5c76d69
Ignore mypy
PierreQuinton Oct 20, 2025
34f4dce
Remove useless comment.
PierreQuinton Oct 20, 2025
a0b7ffc
Change repr
PierreQuinton Oct 20, 2025
f476b29
revert removing `__init__`
PierreQuinton Oct 20, 2025
447d714
Give implementation for pointwise
PierreQuinton Oct 21, 2025
85556a8
Add decorator to handle other functions. Add two examples of such fun…
PierreQuinton Oct 21, 2025
c5f868c
Improve naming.
PierreQuinton Oct 21, 2025
7f5c097
Merge branch 'main' into block-diagonal-tensor
ValerianRey Oct 22, 2025
efa8019
improve
PierreQuinton Oct 22, 2025
e36b3c5
Remove inplace functions from the list of pointwise functions (they s…
PierreQuinton Oct 23, 2025
018a994
Fix `to_dense`
PierreQuinton Oct 23, 2025
e91323c
Verify input.
PierreQuinton Oct 23, 2025
f728153
Make a builder for DSPs and move checks in it. This should always be …
PierreQuinton Oct 23, 2025
b2b5d7a
Implement pointwise and inplace pointwise in `_HANDLED_FUNCTIONS`.
PierreQuinton Oct 23, 2025
2b80788
Move Pointwise functions definitions.
PierreQuinton Oct 23, 2025
85c8e41
Clean filed of DST, remove virtual shape, it is just the shape, and m…
PierreQuinton Oct 23, 2025
123d2ce
Merge branch 'dev-new-engine' into block-diagonal-tensor
ValerianRey Oct 23, 2025
7c0bc45
Use DST for initial jac_output
ValerianRey Oct 23, 2025
8509a33
Add test for to_dense and inplace_pointwise
PierreQuinton Oct 24, 2025
0eed3a3
Revert "Clean filed of DST, remove virtual shape, it is just the shap…
PierreQuinton Oct 24, 2025
55a7cbc
Make `contiguous_data` and `v_to_p` public.
PierreQuinton Oct 24, 2025
5796236
Merge remote-tracking branch 'origin/block-diagonal-tensor' into bloc…
PierreQuinton Oct 24, 2025
0ede4ce
Add linting comment when importing aten
ValerianRey Oct 24, 2025
bb165b0
Specify indexing="ij" in meshgrid call
ValerianRey Oct 24, 2025
6e57a3f
Use tuple for v_indices_grid
ValerianRey Oct 24, 2025
8351af5
Use lists for _POINTWISE_FUNCTIONS and _IN_PLACE_POINTWISE_FUNCTIONS
ValerianRey Oct 24, 2025
278bf24
Sort _POINTWISE_FUNCTIONS and _IN_PLACE_POINTWISE_FUNCTIONS
ValerianRey Oct 24, 2025
d14a1a5
Fix pointwise function override
ValerianRey Oct 24, 2025
17abe7f
Make `contiguous_data` and `v_to_p` public.
PierreQuinton Oct 25, 2025
6bbd702
Use wraps decorator instead of update_wrapper
ValerianRey Oct 26, 2025
fac5f7d
Make densification explicit in tests before assert_close
ValerianRey Oct 26, 2025
68bb2a3
Revert "Make `contiguous_data` and `v_to_p` public."
ValerianRey Oct 26, 2025
3db55c3
Add pow implementation for DST
ValerianRey Oct 26, 2025
90b205c
Remove type hint of __torch_dispatch__ so that mypy stops complaining
ValerianRey Oct 26, 2025
2b36084
Make signature of to_dense match that of Tensor.to_dense.
ValerianRey Oct 26, 2025
6236baa
Change type: ignore so that mypy doesn't complain
ValerianRey Oct 26, 2025
6d28918
Add unsqueeze_default implementation for DST
ValerianRey Oct 26, 2025
f4a436c
Add more info in the print when falling back to dense
ValerianRey Oct 26, 2025
ad4d843
Add implementation for trivial views
ValerianRey Oct 26, 2025
b353ff1
Add expand_default, div_Scalar, and slice_Tensor
ValerianRey Oct 26, 2025
d640266
Add mul_Tensor and transpose_int
ValerianRey Oct 26, 2025
6145e4b
Use diagonal_sparse_tensor in view_default
ValerianRey Oct 26, 2025
444e7cd
Use DiagonalSparseTensor type hint when applicable
ValerianRey Oct 26, 2025
5c859b5
Add debug_info
ValerianRey Oct 26, 2025
960de60
Improve error message in view_default
ValerianRey Oct 26, 2025
e12080e
Add multi-dimensional v_to_p
ValerianRey Oct 26, 2025
a38f907
Rename mean and sum to mean_default and sum_default
ValerianRey Oct 26, 2025
b7a3f53
Fix condition and comment in diagonal_sparse_tensor
ValerianRey Oct 27, 2025
aa3d8a0
Improve implementation of view
ValerianRey Oct 27, 2025
c49411f
Implement `einsum` for `v_to_p: list[int]`
PierreQuinton Oct 27, 2025
6c0e9ec
Factor access to physical shape in view.
PierreQuinton Oct 27, 2025
2f774db
rename `physical_to_virtual` to `p_to_vs`.
PierreQuinton Oct 27, 2025
7344118
Refactor p_to_vs.
PierreQuinton Oct 27, 2025
0884495
Adapt `einsum` to `v_to_p: list[list[int]]`
PierreQuinton Oct 27, 2025
a5fda15
Rename `v_to_p` to `v_to_ps`
PierreQuinton Oct 27, 2025
d03d3a4
Remove diagonal_sparse_tensor, remove to_diagonal_sparse_tensor
ValerianRey Oct 27, 2025
3bf75fb
Remove aten.div.Scalar from _POINTWISE_FUNCTIONS
ValerianRey Oct 27, 2025
35e841a
Fix set union to be able to handle empty list of sets
ValerianRey Oct 27, 2025
2d37031
fix type hint of bmm_default
ValerianRey Oct 27, 2025
7a84940
Add mm_default
ValerianRey Oct 27, 2025
30adf30
Use to_dense when comparing result of einsum in test_einsum
ValerianRey Oct 27, 2025
e6ce46e
Make einsum work with DST only, re-add to_diagonal_sparse_tensor
ValerianRey Oct 27, 2025
5253363
Remove list comprehension
PierreQuinton Oct 27, 2025
75d138c
Add sortin function
PierreQuinton Oct 27, 2025
f33107b
Add test_view2 and new parametrization for test_view
ValerianRey Oct 28, 2025
17047f6
Add test_sort_dst
ValerianRey Oct 28, 2025
6d1cfa5
Revert "Remove list comprehension"
ValerianRey Oct 28, 2025
a5390b7
Rename map to mapping in sort_dst
ValerianRey Oct 28, 2025
2cc95ea
Use empty ps to indicate virtual dim of size one.
ValerianRey Oct 28, 2025
ef5a2b1
Remove outdated todo in einsum
ValerianRey Oct 28, 2025
ea7d4dd
Improve implementation of sort_dst
ValerianRey Oct 28, 2025
8403b11
Restructure packages
ValerianRey Oct 28, 2025
f2b2ef6
Make _HANDLED_FUNCTIONS and implements class attributes of DST
ValerianRey Oct 28, 2025
97213eb
Move pointwise function implementations to the end of the file and av…
ValerianRey Oct 28, 2025
866b4b1
Uniformize name of physical
ValerianRey Oct 28, 2025
b73de79
Improve print when falling back to dense
ValerianRey Oct 28, 2025
e393096
Rename first_sort to encode_by_order and improve its docstring
ValerianRey Oct 28, 2025
7a8b8a4
Improve repr of DST
ValerianRey Oct 28, 2025
28c72ed
Minor formatting fix
ValerianRey Oct 28, 2025
1c1885c
Move _strides to where it is used:
ValerianRey Oct 28, 2025
09f3efc
Minor reformating
ValerianRey Oct 28, 2025
5703fd2
Add check that v_to_ps are correctly encoded
ValerianRey Oct 28, 2025
8a13d49
Fix to_diagonal_sparse_tensor to not create physical with dim of size 1.
ValerianRey Oct 28, 2025
82d4b33
Add function to create a DST without having to care about dims of siz…
ValerianRey Oct 28, 2025
9d21a6c
Remove test_view and rename test_view2 to test_view
ValerianRey Oct 28, 2025
42e8228
Minor reformatting
ValerianRey Oct 28, 2025
33194c0
Move p_to_vs outside of DiagonalSparseTensor and rename it p_to_vs_fr…
ValerianRey Oct 28, 2025
21ca8e3
Add get_groupings, fix_ungrouped_dims, and use it in constructor
ValerianRey Oct 28, 2025
15a615f
Add check of maximal grouping in DST.__init__
ValerianRey Oct 28, 2025
ee8e181
Use make_dst in einsum
ValerianRey Oct 28, 2025
8992d3f
Add possibility to slice dimension of size 1.
ValerianRey Oct 28, 2025
cf4f950
Fix to_dense when a pdims is []
ValerianRey Oct 28, 2025
ebffc3c
Stop creating index tensors on data device. I think it's ok / faster …
ValerianRey Oct 29, 2025
fe84b80
Rename current_product to current_size
ValerianRey Oct 29, 2025
92bcaa6
Add unsquash_pdim
ValerianRey Oct 29, 2025
ceeeea6
Revamp unsquash_dim:
ValerianRey Oct 29, 2025
6ab6b4c
Add possibility to unsquash physical dimensions in view.
ValerianRey Oct 29, 2025
b746af8
Fix get_groupings and add test that failed before this fix
ValerianRey Oct 29, 2025
59d1f2c
Fix dim encoding in transpose_int
ValerianRey Oct 29, 2025
4c48523
Name more variables in __torch_dispatch__
ValerianRey Oct 29, 2025
123ac3c
Remove todos in view_default
ValerianRey Oct 29, 2025
469ddb3
Add support for shape inference in view_default
ValerianRey Oct 29, 2025
bd3e569
Add _unsafe_view_default
ValerianRey Oct 29, 2025
42462bc
Add threshold_backward_default
ValerianRey Oct 29, 2025
19e049d
Add sum_dim_IntList
ValerianRey Oct 29, 2025
519d11b
Add broadcast_tensors_default
ValerianRey Oct 29, 2025
64d0ae8
Fix mul_Tensor
ValerianRey Oct 29, 2025
47f32f8
Add missing default value for keepdim in sum_dim_IntList
ValerianRey Oct 29, 2025
b4eb021
Improve print when falling back to dense
ValerianRey Oct 29, 2025
909bcb7
Add mul_Scalar
ValerianRey Oct 29, 2025
c72f99e
Fix mul_Tensor to be able to handle non-tensor input
ValerianRey Oct 29, 2025
4b3593b
Make test_einsum parametrizable
ValerianRey Oct 29, 2025
2052d81
Add failing example to test_einsum
ValerianRey Oct 29, 2025
a1ccad5
Add even simpler failing example in test_einsum
ValerianRey Oct 29, 2025
2675d8e
Add assertion about result being DST in test_einsum
ValerianRey Oct 29, 2025
3121b6a
Add a test to get_groupings. May want to add the same to `test_fix_un…
PierreQuinton Oct 29, 2025
048a1e2
Merge remote-tracking branch 'origin/block-diagonal-tensor' into bloc…
PierreQuinton Oct 29, 2025
5a2a7a0
Add `strides_from_p_dims_and_p_shape`
PierreQuinton Oct 29, 2025
4b83644
Add `merge_strides`
PierreQuinton Oct 29, 2025
7e8439a
Add `stride_to_shape`
PierreQuinton Oct 29, 2025
c7047e1
Add `to_target_physical_strides`
PierreQuinton Oct 29, 2025
c6e3fd9
Add new_implementation idea in einsum.
PierreQuinton Oct 29, 2025
cecd69b
Merge branch 'dev-new-engine' into block-diagonal-tensor
ValerianRey Oct 30, 2025
d2e53a3
Always use randn_ in test
ValerianRey Oct 30, 2025
16e7c1c
Fix order of sorting in merg_strides
PierreQuinton Oct 30, 2025
2c23dbe
Merge remote-tracking branch 'origin/block-diagonal-tensor' into bloc…
PierreQuinton Oct 30, 2025
16e6165
Add strides_v2
ValerianRey Oct 30, 2025
59cf10b
Add more failing parametrizations to test_get_groupings and test_fix_…
ValerianRey Oct 30, 2025
e0dc1a7
Add test_concatenate
ValerianRey Oct 30, 2025
3dffd1e
Add strides_to_pdims
ValerianRey Oct 30, 2025
48387ab
Add (passing) test_to_dense2 to test to_dense when the tensor has a v…
ValerianRey Oct 31, 2025
f87ecbb
Use new strides in to_dense
ValerianRey Oct 31, 2025
b883f20
Simplify to_dense
ValerianRey Oct 31, 2025
7044e37
Remove unused variable dims in the loop of to_dense
ValerianRey Oct 31, 2025
384d550
Pre-compute strides
ValerianRey Oct 31, 2025
18044dd
Replace for-loop with for comprehension to create v_indices_grid, and…
ValerianRey Oct 31, 2025
b0a0e7a
Remove for-loop in computation of v_indices_grid in to_dense
ValerianRey Oct 31, 2025
456adf1
Add internal strides to debug_info
ValerianRey Oct 31, 2025
dc37d29
Simplify creation of strides in to_dense
ValerianRey Oct 31, 2025
a9e298a
Remove torch. prefix when possible
ValerianRey Oct 31, 2025
ac2a2c5
Create strides tensor in constructor
ValerianRey Oct 31, 2025
cfba7e0
Remove comments about the device of the indices tensors. We actually …
ValerianRey Oct 31, 2025
2c94488
Replace self.physical.device and self.physical.dtype by self.device a…
ValerianRey Oct 31, 2025
9c1ad5b
Move unwrap_to_dense out of DiagonalSparseTensor
ValerianRey Oct 31, 2025
7252757
Extract print_fallback
ValerianRey Oct 31, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/torchjd/sparse/_diagonal_sparse_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,10 @@ def strides_from_p_dims_and_p_shape(p_dims: list[int], physical_shape: list[int]
]


def merge_strides(strides: list[list[int]]) -> list[int]:
return sorted({s for stride in strides for s in stride})


def p_to_vs_from_v_to_ps(v_to_ps: list[list[int]]) -> list[list[tuple[int, int]]]:
"""
A physical dimension is mapped to a list of couples of the form
Expand Down