Skip to content

Replace RecursiveApply functions with AutoBroadcaster wrappers#2417

Merged
dennisYatunin merged 1 commit intomainfrom
dy/math_wrapper
May 1, 2026
Merged

Replace RecursiveApply functions with AutoBroadcaster wrappers#2417
dennisYatunin merged 1 commit intomainfrom
dy/math_wrapper

Conversation

@dennisYatunin
Copy link
Copy Markdown
Member

@dennisYatunin dennisYatunin commented Jan 10, 2026

Purpose

This 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 external libraries like ClimaAtmos, we can now use a wrapper for Tuples and NamedTuples that behaves the same way when passed to standard math functions. The number of places where this interface needs to be explicitly used is quite small, making it possible to decouple iteration over field variables from most of ClimaCore's internals.

Content

  • Introduce the AutoBroadcaster wrapper, which triggers nested_broadcast (an analogue of the old RecursiveApply.rmap) when passed to any common math function or constructor
    • Replace all RecursiveApply functions in ClimaCore's source code with standard math operations
    • Make it possible to add support for any iterator type T by defining the method is_auto_broadcastable(::T) = true, as long as T is compatible with UnrolledUtilities
  • Preserve backward-compatibility with original RecursiveApply implementation
    • Add a RecursiveApply module that defines the same functions as the old version, but as simple wrappers around add_auto_broadcasters and drop_auto_broadcasters
    • Define an auto_broadcasted function that automatically introduces calls to either add_auto_broadcasters or drop_auto_broadcasters in Broadcasted expressions, allowing ClimaCore operators to be used in the same expressions as user-defined functions that are incompatible with AutoBroadcasters (e.g., functions that dispatch on non-AutoBroadcaster types)
  • Unify broadcasting across DataLayouts, Fields, and Operators
    • The broadcastable form of any AbstractData or Field has add_auto_broadcasters applied to its values
    • Calling Base.broadcasted triggersauto_broadcasted instead of the standard Broadcasted constructor
    • The result of eltype is the output of unsafe_eltype, which acts as an unrolled version of Base.Broadcast.combine_eltypes
    • The result of similar uses the output of safe_eltype, which can identify the first point during inference where a non-concrete type appears in a broadcast expression (essentially a simplified and corrected version of the old call_with_first)
  • Update the implementation of new to avoid segfaults and other errors when dealing with DataType fields
  • Refactor the deformation_flow.jl example by removing all loops over field variables, adding several missing cases, and clarifying precisely what needs to be tested

  • 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 10 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

Comment thread src/Utilities/auto_broadcaster.jl
Comment thread src/Utilities/auto_broadcaster.jl Outdated
Comment thread src/Utilities/auto_broadcaster.jl Outdated
Comment thread src/Utilities/auto_broadcaster.jl Outdated
Comment thread src/Utilities/Utilities.jl Outdated
Comment thread src/Utilities/Utilities.jl Outdated
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