Skip to content

Replace RecursiveApply functions with AutoBroadcaster wrappers#2417

Open
dennisYatunin wants to merge 1 commit intomainfrom
dy/math_wrapper
Open

Replace RecursiveApply functions with AutoBroadcaster wrappers#2417
dennisYatunin wants to merge 1 commit intomainfrom
dy/math_wrapper

Conversation

@dennisYatunin
Copy link
Copy Markdown
Member

@dennisYatunin dennisYatunin commented Jan 10, 2026

Replaces the RecursiveApply module with a new interface that follows the opposite design pattern: instead of defining functions for Tuples and NamedTuples that need to be called in Fields, Operators, MatrixFields, and sometimes even in ClimaAtmos, we can now use a wrapper for Tuples and NamedTuples that behaves the same way when passed to any standard math function. This should make it easy to decouple iteration over field variables from most of ClimaCore's internals.

The new interface can be extended to any iterator type T by defining the method is_auto_broadcastable(::T) = true, as long as T is compatible with UnrolledUtilities.jl. The @auto_broadcaster_args annotation provides a convenient shorthand for extending the new interface to arbitrary functions of Numbers or AbstractArrays, automatically broadcasting those functions over Tuples, NamedTuples, and AutoBroadcasterss. This annotation can also be used to add methods for any other type T, with the resulting behavior depending on whether is_auto_broadcastable is true for that type.

  • Code follows the style guidelines OR N/A.
  • Unit tests are included OR N/A.
  • Code is exercised in an integration test OR N/A.
  • Documentation has been added/updated OR N/A.

@dennisYatunin dennisYatunin force-pushed the dy/math_wrapper branch 9 times, most recently from cc7a736 to 0b0cbe1 Compare January 16, 2026 03:34
@dennisYatunin dennisYatunin force-pushed the dy/math_wrapper branch 2 times, most recently from e603f47 to 5430722 Compare January 21, 2026 17:33
@dennisYatunin dennisYatunin changed the title Replace RecursiveApply interface with MathWrapper Replace RecursiveApply interface with MathMapper Jan 21, 2026
@dennisYatunin dennisYatunin force-pushed the dy/math_wrapper branch 14 times, most recently from 30377bd to 92cf627 Compare January 24, 2026 02:30
@dennisYatunin dennisYatunin force-pushed the dy/math_wrapper branch 11 times, most recently from 2bd7d7c to 52d17a0 Compare February 3, 2026 04:41
@dennisYatunin dennisYatunin changed the title Replace RecursiveApply interface with MathMapper Replace RecursiveApply functions with AutoBroadcaster wrappers Feb 3, 2026
@dennisYatunin dennisYatunin force-pushed the dy/math_wrapper branch 12 times, most recently from 8553df5 to 264d4d2 Compare February 5, 2026 01:44
Copy link
Copy Markdown
Member

@nefrathenrici nefrathenrici left a comment

Choose a reason for hiding this comment

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

This looks good to me. It's nice to get rid of RecursiveApply. I didn't have any significant comments.

Comment thread ext/cuda/operators_spectral_element.jl Outdated
Comment thread src/MatrixFields/MatrixFields.jl Outdated
Copy link
Copy Markdown
Member

@imreddyTeja imreddyTeja left a comment

Choose a reason for hiding this comment

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

I left some comments and questions, but this looks mostly good to me

which it is true, while leaving values for which it is false unmodified. Can
also be passed an iterator's type to infer the result type for such an iterator.
"""
add_auto_broadcasters(itr) =
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I don't think we'd ever see structs nested that much, but should this have the recursion limit removed?

wrappers no longer corresponds to a guaranteed error, it is returned instead of
the original `bc`. Otherwise, `bc` is returned without any modifications.
"""
function allow_auto_broadcasting(bc::Base.Broadcast.Broadcasted)
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

has this function been checked for allocations or dynamic function invokations?

AutoBroadcaster(merge(unrolled_map(unwrap, args)...))

for f in (:axes, :size, :firstindex, :lastindex)
@eval Base.$f(x::AutoBroadcaster) = $f(unwrap(x))
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Could this be combined with the loop on line 145?

) where {F, O} =
unrolled_mapreduce(f, op, unwrap(arg), unrolled_map(unwrap, args)...; init...)
Base.map(f::F, arg::AutoBroadcaster, args::AutoBroadcaster...) where {F} =
AutoBroadcaster(unrolled_map(f, unwrap(arg), unrolled_map(unwrap, args)...))
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

How does this work if the arg is a AutoBroadcaster wrapping a tuple of AutoBroadcasters? Wouldn't it need to recurse with Base.map instead?

ismutabletype(T) || unrolled_all(no_immutable_types, fieldtypes(T))

drop_immutable_types(not_a_type) = not_a_type
drop_immutable_types(::Type{<:Type}) = DataType
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

What is this doing?


no_immutable_types(::Type{<:Type}) = false
no_immutable_types(::Type{T}) where {T} =
ismutabletype(T) || unrolled_all(no_immutable_types, fieldtypes(T))
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This function could use a comment or docstring. How does unrolled_all(no_immutable_types, fieldtypes(T)) imply the type is immutible

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