Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rework cat implementation #18

Merged
merged 28 commits into from
Mar 10, 2025
Merged

Rework cat implementation #18

merged 28 commits into from
Mar 10, 2025

Conversation

lkdvos
Copy link

@lkdvos lkdvos commented Feb 4, 2025

I cooked up a minimal version of the parts of a small adaptation to the Base implementation of cat with more well-defined entry-points for customization, and which should be friendly on both inference and invalidations because the dispatching is happening in steps.

Let me know how you feel about this and we can think about the most convenient way of integrating this into the rest of the framework.

Copy link

codecov bot commented Feb 4, 2025

Codecov Report

Attention: Patch coverage is 78.94737% with 8 lines in your changes missing coverage. Please review.

Project coverage is 72.72%. Comparing base (f691e77) to head (5084093).
Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
src/concatenate.jl 80.55% 7 Missing ⚠️
src/traits.jl 0.00% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main      #18      +/-   ##
==========================================
+ Coverage   72.44%   72.72%   +0.28%     
==========================================
  Files           9       10       +1     
  Lines         323      308      -15     
==========================================
- Hits          234      224      -10     
+ Misses         89       84       -5     
Flag Coverage Δ
docs 38.88% <0.00%> (+1.70%) ⬆️

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
Copy link
Member

I think it is ok, I see that you tried to stick closer to the implementation of Base.cat. However, I think I would prefer a strategy more like the design of Base.Broadcast.Broadcasted, which as you know stores the expression and a style based on combining the styles of the arguments. Having to overload cat_similar(::Type{T}, shape, args...) to get the destination object isn't ideal since you still have the issue of having many potential combinations of arguments to deal with.

I think we can have some compromise between the two approaches, for example:

function cat(dims, args...)
  # `instantiate` reduces over the `args` and determines a common `CatStyle`,
  # returning a new `Catted` with the `CatStyle` stored.
  return copy(instantiate(Catted(dims, args...)))
end

function Base.copy(c::Catted)
  dest = similar(c)
  copyto!(dest, c)
  return dest
end

function Base.copyto!(dest::AbstractArray, c::Catted)
  cat!(dest, axes(c), c.dims, c.args...)
  return dest
end

function cat!(dest, shape, catdims, args...)
  # Implementation of `cat!`.
end

@lkdvos
Copy link
Author

lkdvos commented Feb 5, 2025

I initially started with a design similar to Base.Broadcast, but quickly realized that it looks like a huge overkill for this specific situation.
Given that we really only want need to make sure the destination has the right type, my current approach is the minimal amount of changes (and thus code to test, document, understand and maintain) I could come up with to make that happen.
In particular, I like that this could reuse the interfaces for defining an implementation of @interface cat_similar, so we don't have to introduce a new style, new combination rules, etc.

I agree that sure, the Broadcast approach is more general, more customizable and more flexible, but it's also more complicated and I am not convinced we really need that. There's this quote about "every line of code you write is a potential bug", which I think could be applicable here :).

Concretely, I would suggest we avoid the Catted object and keep the current structure, but have a default implementation for cat_similar(::Type{T}, shape, args...) = cat_similar(interface(args...), T, shape, args...). This keeps the simplicity and still has the flexibility?

@mtfishman
Copy link
Member

Concretely, I would suggest we avoid the Catted object and keep the current structure, but have a default implementation for cat_similar(::Type{T}, shape, args...) = cat_similar(interface(args...), T, shape, args...). This keeps the simplicity and still has the flexibility?

I'm a bit concerned about reusing interface for this purpose, does that really output the correct type or give the right level of customization in general (say when mixing different kinds of sparse types)?

@lkdvos
Copy link
Author

lkdvos commented Feb 5, 2025

It's possible, but then we can still specialize on the types. Also, since we currently don't have that many different sparse types in place, I would much rather have a simple solution that works for now, which we can fix later if need be, instead of a more complicated one that may account for future options, but may also still have to be fixed later because we cannot foresee all possibilities without actual examples.

@mtfishman
Copy link
Member

I still personally prefer a solution based around a Catted object (like the current implementation in the package, though of course it doesn't have to follow that code exactly), I don't think it is so complicated to implement and is a way to help organize the code. For now we can forgo defining CatStyle to keep things simpler, for example use interface as a customization point as you suggested (Catted could store the result of interface(args...), for example).

@lkdvos
Copy link
Author

lkdvos commented Feb 11, 2025

I reworked the implementation to include a Concatenated object (this somehow felt like a better name than Catted, but I'm happy to switch it out). Let me know how you feel about the general implementation here, and I'll try to hook it into the rest of the machinery?

@mtfishman
Copy link
Member

mtfishman commented Feb 11, 2025

I like the new names, Concatenated definitely looks better than Catted.

What do you have in mind for how to incorporate this into the rest of the package? I'm curious what you think of: #12 (comment). It seems like those broader questions around what this package should be and how opinionated it should be has implications for what we decide to do with these cat methods. If we go the route of having this package be very unopinionated and not really having any implementations, it might then make sense for this PR to be put in a separate package Concatenated.jl, more analogous to https://github.com/ITensor/MapBroadcast.jl, and then other packages like SparseArraysBase.jl can opt-in to that implementation of cat, say by forwarding Base.cat/Base._cat to Concatenated.concatenate or @interface ConcatenatedArrayInterface() Base.cat (which could be defined as Concatenated.concatenate).

@lkdvos
Copy link
Author

lkdvos commented Feb 11, 2025

I think I would lean towards having defaults that aren't opinionated, while providing tools for customization.
It seems fair to me to have this machinery as an opt-in option, indeed by having Base._cat(dims, args...) = concatenate(args...) and then using the concatenation-specific implementation. The annoying part is mostly that there is no way to fix the Base.cat(::Array1, ::Array2) ambiguities that arise from definitions like cat(::Array1, ::AbstractArray), except for manually dispatching the offending signatures to concatenate.
In some sense, it would almost be more convenient to start using concatenate instead of cat, but that might be too much to ask too...

I'm not very enthusiastic about splitting this off into a separate package. This has the drawback that it would then require us to also add the logic of having interfaces and interface promotion rules there, or simply depend on this package anyways. I fear that if we need to have Interface, CatStyle, BroadcastStyle, and support rules for combining them, that becomes a bit much, and if we need to depend on this package anyways, then I don't think there's that much to gain from splitting it off.

Another somewhat unfortunate factor is that these kinds of solutions don't follow the same format as the other functions, and it's a bit hard to try and fit it into the @derive signature. This doesn't need to be a problem, I'm definitely okay with just manually writing Base._cat(dims, args...) = DerivableInterfaces.concatenate(args...; dims) (In fact, I even prefer the more verbose way), since this makes it very clear what is going on and where.

@mtfishman
Copy link
Member

I think I would lean towards having defaults that aren't opinionated, while providing tools for customization. It seems fair to me to have this machinery as an opt-in option, indeed by having Base._cat(dims, args...) = concatenate(args...) and then using the concatenation-specific implementation. The annoying part is mostly that there is no way to fix the Base.cat(::Array1, ::Array2) ambiguities that arise from definitions like cat(::Array1, ::AbstractArray), except for manually dispatching the offending signatures to concatenate. In some sense, it would almost be more convenient to start using concatenate instead of cat, but that might be too much to ask too...

Yeah, I think that would be a bit too much (it would be unfortunate for someone to call Base.cat on a SparseArrayDOK and have it use a bad or broken Base implementation...), but we can think of it as a proposal for a potential new version of Base.cat, like how ArrayLayouts.jl started out as a proposed rewrite of parts of LinearAlgebra.jl.

I'm not very enthusiastic about splitting this off into a separate package. This has the drawback that it would then require us to also add the logic of having interfaces and interface promotion rules there, or simply depend on this package anyways. I fear that if we need to have Interface, CatStyle, BroadcastStyle, and support rules for combining them, that becomes a bit much, and if we need to depend on this package anyways, then I don't think there's that much to gain from splitting it off.

Right, for now I was picturing that if we moved it to a separate package we could depend on this package and the interface mechanism for inferring the output interface/type. I think the gain is conceptual/organizational, I just don't want this package growing into a "Base alternative", which you already identified it was doing a bit too much already.

Another somewhat unfortunate factor is that these kinds of solutions don't follow the same format as the other functions, and it's a bit hard to try and fit it into the @derive signature. This doesn't need to be a problem, I'm definitely okay with just manually writing Base._cat(dims, args...) = DerivableInterfaces.concatenate(args...; dims) (In fact, I even prefer the more verbose way), since this makes it very clear what is going on and where.

As you say, I think that is just the nature of the fact that Base.cat isn't good to overload, and is unavoidable right now (unless we define an @interface version of Base._cat but I'm not a big fan of that approach...).

I'm ok keeping it in this package right now, but lets keep it in mind going forward. I would want to have some "philosophy" or guiding principle for what goes in this package vs. in a separate package. Maybe the principle is that if it implements functionality that is close to functionality in Base but provides better customization points based around interface, it can go in this package.

@mtfishman
Copy link
Member

@lkdvos this looks good to me, what are the next steps here?

@lkdvos
Copy link
Author

lkdvos commented Feb 12, 2025

Since this one is breaking, I would suggest either merging but not yet tagging, or slightly holding off on merging, because maybe we can simultaneously alter the AbstractArrayInterface having too many different interpretations, as mentioned in #12 .
Concretely, I would suggest changing the following breaking changes before tagging a release:

  • remove AbstractArrayOps
  • new concatenate API
  • AbstractArrayInterface rework

@mtfishman
Copy link
Member

Sure, that sounds reasonable to do all of those in a single breaking release. So, up to you if you want to merge this and then work on those next, or keep this unmerged and building off of it (I think I prefer not merging this yet since if there are non-breaking things we want to fix in the meantime it would make that easier).

@lkdvos
Copy link
Author

lkdvos commented Feb 12, 2025

I'm okay with keeping this open, the other changes should not depend on this anyways

@lkdvos
Copy link
Author

lkdvos commented Mar 8, 2025

If tests turn green, I think this should be ready to release.
This is (slightly) breaking, and downstream packages will have to:

  • implement Concatenate.cat instead of the interfaced Base.cat
  • implement DerivableInterfaces.zero! instead of ArrayLayouts.zero!

@lkdvos
Copy link
Author

lkdvos commented Mar 9, 2025

Ready to go from my side.

@mtfishman
Copy link
Member

Can you tell what's going on with the downstream tests? I would have hoped bumping the version would mean the downstream packages fail during the installation step, while it seems like they are downgrading to some broken state. But maybe what's happening is that the downstream tests aren't all testing direct dependencies, so then they don't have compat entries for DerivableInterfaces. I'm guessing it is fine from the endpoint of a user unless they force DerivableInterfaces to use v0.4, I just want to double check if you see anything we should try to change.

f = eltype(a) <: Number ? Returns(zero(eltype(a))) : zero!
return @interface interface map!(f, a, a)
end
# TODO: should this be recursive? `map!(zero!, A, A)` might also work?
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The case I can think of where we might want it to be recursive is for the blocks/storage of BlockSparseArray, which as you know is a sparse array of arrays. Do you know how this definition affects that case? Maybe there we don't ever call zero!(blocks(a)) and we just define zero!(a), i.e. define it directly on the block sparse array?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, just to clarify, is the idea now that if types want to overload ArrayLayouts.zero!, they can do some "manually", for example by defining it as DerivableInterfaces.zero!?

Copy link
Author

@lkdvos lkdvos Mar 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be fair, for BlockSparseArrays I would expect a custom implementation anyways, since we probably want to dispatch that to zero!(blocks(a)), and we would also need a custom implementation for SparsArrayDOK anyways, since we don't actually want to fill anything at all, but rather delete?

And yes, I think the idea was to decouple this entirely from ArrayLayouts? I don't have any special opinions about this really, I just want a function to overload when necessary

Copy link
Member

@mtfishman mtfishman Mar 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be fair, for BlockSparseArrays I would expect a custom implementation anyways, since we probably want to dispatch that to zero!(blocks(a)), and we would also need a custom implementation for SparsArrayDOK anyways, since we don't actually want to fill anything at all, but rather delete?

That sounds right, I just wanted to make sure I was following along with what the plan is there so I knew how to review this change, and also probe whether we want zero! to be recursive. My guess is that we would want it to be, though probably we can hold off on that for now since as you say our use cases require custom implementations anyway.

And yes, I think the idea was to decouple this entirely from ArrayLayouts? I don't have any special opinions about this really, I just want a function to overload when necessary

That's what I had in mind and I think what we had discussed, I just wanted to make sure we had a plan there that we both agreed on. (I think more generally, we basically want to leave it up to types to overload other interfaces like the ArrayLayouts.jl interface, the broadcast interface, etc. and not have that logic in DerivableInterfaces.jl.)

@lkdvos
Copy link
Author

lkdvos commented Mar 10, 2025

Can you tell what's going on with the downstream tests? I would have hoped bumping the version would mean the downstream packages fail during the installation step, while it seems like they are downgrading to some broken state. But maybe what's happening is that the downstream tests aren't all testing direct dependencies, so then they don't have compat entries for DerivableInterfaces. I'm guessing it is fine from the endpoint of a user unless they force DerivableInterfaces to use v0.4, I just want to double check if you see anything we should try to change.

Actually, I think Pkg is just being smarter than we are: NamedDimsArrays used to depend on Derive instead of DerivableInterfaces, so it did not have a compat entry for DerivableInterfaces. Therefore Pkg just happily selects that version, which then does not even use DerivableInterfaces at all. Somehow there is still an error in versions there, but that seems like it's a compat issue with BlockSparseArrays and an old version of NamedDimsArrays.

@mtfishman
Copy link
Member

Actually, I think Pkg is just being smarter than we are: NamedDimsArrays used to depend on Derive instead of DerivableInterfaces, so it did not have a compat entry for DerivableInterfaces. Therefore Pkg just happily selects that version, which then does not even use DerivableInterfaces at all. Somehow there is still an error in versions there, but that seems like it's a compat issue with BlockSparseArrays and an old version of NamedDimsArrays.

Hmm ok, I guess that is a bit unfortunate since it would be nice to have the tests fail "properly" if we make breaking changes to DerivableInterfaces.jl, rather than force a downgrade to older package versions. Would it make sense to yank those older versions of those packages that depend on Derive.jl? I'd be ok going a bit "scorched earth" on the ITensorRegistry, we don't really need those older versions and packages to be installable at this point since so much has changed and improved since then.

@lkdvos lkdvos merged commit 14cb954 into ITensor:main Mar 10, 2025
12 of 16 checks passed
@lkdvos lkdvos deleted the cat-2 branch March 10, 2025 13:29
@lkdvos
Copy link
Author

lkdvos commented Mar 10, 2025

Yeah, I think removing these from the registry is more than fair, given that this is not production-ready anyways I don't see the need to be that backwards compatible

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants