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

Store state node IDs in parameters #2086

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Store state node IDs in parameters #2086

wants to merge 1 commit into from

Conversation

visr
Copy link
Member

@visr visr commented Feb 21, 2025

So we can remove id_from_state_index.

So we can remove `id_from_state_index`.
@visr
Copy link
Member Author

visr commented Feb 21, 2025

These are scraps from #2077. Thinking about this, and having to deal a bit more with huge stacktraces from the large ComponentArrays types, made me consider another approach. Better in a separate PR, but dumping thoughts here. We could store the UnitRange{Int64} of each state component in the parameters, and use view(u, p.is_pump) instead of u.pump. Storing the resulting SubArray like in StateVector from #2077 can maybe also be done, but is more complex because you'd have to support Duals and GradientTracer as well as Float64.

I'm hoping this can also help with latency (#1942). Since ComponentArrays puts the full Axes into its type signature, it will specialize differently on each model. Because if the state length or component lengths differ, it needs to compile separate specializations.

# uses MethodAnalysis, f(x) can be anything
julia> specializations = methodinstances(f)
2-element Vector{Core.MethodInstance}:
 MethodInstance for f(::ComponentVector{Int64, Vector{Int64}, Tuple{Axis{(a = ViewAxis(1:3, Shaped1DAxis((3,))), b = ViewAxis(4:5, Shaped1DAxis((2,))))}}})
 MethodInstance for f(::ComponentVector{Int64, Vector{Int64}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = ViewAxis(3:5, Shaped1DAxis((3,))))}}})

Here the two ComponentVectors only differ slightly:

julia> using ComponentArrays
julia> a = ComponentVector(; a = [1,2,3], b = [4,5])
julia> b = ComponentVector(; a = [1,2], b = [3,4,5])
julia> typeof(a)==typeof(b)
false

Many methods that we call are likely to already be precompiled for simple Vector{Float64}, but none will be for our ComponentVector state (or StateVector).

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.

1 participant