Skip to content

Conversation

@AayushSabharwal
Copy link
Member

@AayushSabharwal AayushSabharwal commented Sep 16, 2025

Requires JuliaSymbolics/SymbolicUtils.jl#795

I haven't pushed test changes since they could need multiple rounds of updates, bloating the diff and making it harder to rebase.

Copy link
Contributor

@SebastianM-C SebastianM-C left a comment

Choose a reason for hiding this comment

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

Thanks for the update, I left some questions

Comment on lines -7 to +9
using Symbolics.SymbolicUtils
using Symbolics: SymbolicT, VartypeT
using SymbolicUtils
using SymbolicUtils: symtype, unwrap
Copy link
Contributor

Choose a reason for hiding this comment

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

I was doing using Symbolics.SymbolicUtils to laod SymbolicUtils from the Symbolics dependency. I think that was done so that we don't have both Symbolics & SymbolicUtils as extension triggers.
Would it make sense to split the extension or to have 2 triggers given that the SU part is now more significant?

Copy link
Member Author

Choose a reason for hiding this comment

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

Lux is the trigger, not Symbolics/SymbolicUtils. Dependencies of the parent package (Symbolics) are directly available to extensions without requiring a transitive import or additional triggers.

function SymbolicUtils.promote_shape(f::LuxStatelessApplicator, x::SymbolicUtils.ShapeT, ::SymbolicUtils.ShapeT)
@nospecialize x ps
model = __get_model(f)
x = x::SymbolicUtils.ShapeVecT
Copy link
Contributor

Choose a reason for hiding this comment

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

what does this do? 😅
Why don't we have x::SymbolicUtils.ShapeVecT in the function signature?

Copy link
Member Author

Choose a reason for hiding this comment

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

promote_shape takes the operation f along with the shapes of its arguments, and returns the shape of the result. The arguments are all annotated ::ShapeT because that is the type of the shape field in BasicSymbolic. Any call to promote_shape will at best infer ::ShapeT for the shape arguments. Annotating the argument with ::ShapeVecT would mean that the compiler doesn't know at compile time whether it can dispatch to this method or not. @nospecialize tells the compiler to not compile methods for more specific types of x and ps. This helps avoid dynamic dispatch.

@AayushSabharwal AayushSabharwal force-pushed the as/su-v4 branch 2 times, most recently from b71cdc2 to d523711 Compare November 4, 2025 09:07
@AayushSabharwal AayushSabharwal merged commit 3ec4b9f into master Nov 4, 2025
15 of 19 checks passed
@AayushSabharwal AayushSabharwal deleted the as/su-v4 branch November 4, 2025 13:09
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.

3 participants