More customization points in Concatenate

[mtfishman]

This is a followup to #18. In #18, the logic for determining the output shape is based on the Base logic, which is very size-oriented and doesn't allow for customization points for custom axes like blocked axes. This PR is meant to address that, which helps when implementing concatenation in BlockSparseArrays.jl (ITensor/BlockSparseArrays.jl#84).

[codecov]

Codecov Report

Attention: Patch coverage is 71.42857% with 18 lines in your changes missing coverage. Please review.

Project coverage is 71.67%. Comparing base (14cb954) to head (9453585).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/concatenate.jl	72.88%	16 Missing ⚠️
...lockArraysExt/DerivableInterfacesBlockArraysExt.jl	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #32      +/-   ##
==========================================
- Coverage   72.72%   71.67%   -1.06%     
==========================================
  Files          10       11       +1     
  Lines         308      353      +45     
==========================================
+ Hits          224      253      +29     
- Misses         84      100      +16     

Flag	Coverage Δ
docs	32.76% <0.00%> (-6.13%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mtfishman]

@lkdvos let me know what you think. Here's a summary of the changes:

1. I'm now computing the axes when Concatenated is constructed and storing it so it doesn't have to be recomputed multiple times (using a new cat_axes function).
2. Computing the axes now uses a custom implementation which ultimately calls down to a binary cat_axis(r1::AbstractUnitRange, r2::AbstractUnitRange) function, which can be customized for range types that have extra structure (like BlockedOneTo, which is handled in a new package extension DerivableInterfacesBlockArraysExt, and will allow defining how GradedOneTo should be combined).
3. I've copied over some of the Base implementation so now we don't call out to private Base functions (to make the code more future-proof in case they refactor the internal implementation, and also it could help if we want to make the implementation more generic in some way).

[lkdvos]

Some general comments:

I'm not sure if it makes sense to include the axes in the type, I might be missing something, but it seems like this is only ever really called once: when you initialize the destination. Since this is not really an expensive call anyways, do we need to cache that? I'm okay with either, just not sure how much that changes.

If we decide to keep it, I do think we probably want an inner constructor where you can pass in the axes yourself. If the only inner constructor computes this, you have to recompute it when you do conversions etc. (also, currently there is no way to not compute it, so the isnothing check in the function is not really doing anything)

[mtfishman]

I'm not sure if it makes sense to include the axes in the type, I might be missing something, but it seems like this is only ever really called once: when you initialize the destination. Since this is not really an expensive call anyways, do we need to cache that? I'm okay with either, just not sure how much that changes.

The reason why I added it is because in the current implementation I define size(::Concatenated) and ndims(::Concatenated) in terms of axes(::Concatenated), and the way the code is written now, size is also called, so it avoids building the axes twice (though that couples back to our discussion about how we should determine the size in the implementation, i.e. if we should use the Base implementation instead). It doesn't seem like a stretch that we might refactor the code where some combination of size, ndims, and axes are called multiple times and I would want those calls to be cheap (and they could become slightly nontrivial for blocked ranges, though of course I doubt it would ever be a bottleneck).

If we decide to keep it, I do think we probably want an inner constructor where you can pass in the axes yourself. If the only inner constructor computes this, you have to recompute it when you do conversions etc. (also, currently there is no way to not compute it, so the isnothing check in the function is not really doing anything)

That's a good point, I'll change that.

[mtfishman]

Looks like ITensor/ITensorRegistry#71 worked and now the downstream tests "fail" in a better way (i.e. the downstream packages can't be installed, so the tests catch that is the case and pass).

[mtfishman]

@lkdvos I think I've addressed most of your comments, let me know what you think.

[lkdvos]

The reason why I added it is because in the current implementation I define size(::Concatenated) and ndims(::Concatenated) in terms of axes(::Concatenated), and the way the code is written now, size is also called, so it avoids building the axes twice (though that couples back to our discussion about how we should determine the size in the implementation, i.e. if we should use the Base implementation instead).

This is true, but I would guess that in the blockaxes case you would actually not go through the default implementation in terms of __cat! and friends anyways, so you would not end up calling that twice. I'm definitely fine with leaving this as is, for sure now that you already added it since it does not hurt of course. It's mostly just me remembering the "don't cache unless necessary" argument, and since . If size really shows up in the profiler, it might make more sense to actually provide a better implementation of that function instead.

[mtfishman]

This is true, but I would guess that in the blockaxes case you would actually not go through the default implementation in terms of __cat! and friends anyways, so you would not end up calling that twice. I'm definitely fine with leaving this as is, for sure now that you already added it since it does not hurt of course. It's mostly just me remembering the "don't cache unless necessary" argument, and since . If size really shows up in the profiler, it might make more sense to actually provide a better implementation of that function instead.

Those are fair points. Partially this is following the lead of Broadcast, which also stores the axes (though only after the broadcast expression gets instantiated), and also a conceptual point that Concatenated should store information about what you need to concatenate the input arguments, and the axes are important non-trivial information you need for that. I'm definitely not in favor of storing arbitrary cache information in structs, especially when the information is implicitly defined from other fields.