Skip to content

Import ChainRules for broadcast(f, As...) for Array{<:IEEEFloat} #737

@yebai

Description

@yebai

Mooncake's default broadcast(f, As...) rules are generally fast, but they can be slow when f is computationally lightweight.

This has been observed in several settings, eg, #521, #249

One might look into the possibility of importing ChainRules rules for broadcast via @from_chainrules for specific input types like Array{<:IEEEFloat}. See below for an example how this might be done

@eval @from_rrule(
DefaultCtx,
Tuple{
typeof(Broadcast.broadcasted),
typeof($f_fast),
Union{IEEEFloat,Array{<:IEEEFloat}},
},
)

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions