We have rules for mean , except for <code class="notra

Ideally it would probably share code, have a function which for <code class="notransla

Add rule for `mean(f, x; dims)` about chainrules.jl HOT 3 OPEN

juliadiff commented on May 26, 2024 1

Add rule for `mean(f, x; dims)`

from chainrules.jl.

Comments (3)

mcabbott commented on May 26, 2024 1

Ideally it would probably share code, have a function which for mean gets scale=1/size(...) or something.

Xref #529 which is trying to re-work that rule.

from chainrules.jl.

mzgubic commented on May 26, 2024

The implementation could look a bit like sum(f, xs):

ChainRules.jl/src/rulesets/Base/mapreduce.jl

Lines 69 to 94 in ce78d3d

 function rrule( 

 config::RuleConfig{>:HasReverseMode}, ::typeof(sum), f, xs::AbstractArray; dims=: 

 ) 

 fx_and_pullbacks = map(x->rrule_via_ad(config, f, x), xs) 

 y = sum(first, fx_and_pullbacks; dims=dims) 

 pullbacks = last.(fx_and_pullbacks) 

 project = ProjectTo(xs) 

 function sum_pullback(ȳ) 

 call(f, x) = f(x) 

 # if dims is :, then need only left-handed only broadcast 

 broadcast_ȳ = dims isa Colon ? (ȳ,) : ȳ 

 f̄_and_x̄s = call.(pullbacks, broadcast_ȳ) 

 # no point thunking as most of work is in f̄_and_x̄s which we need to compute for both 

 f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f 

 NoTangent() 

 else 

 sum(first, f̄_and_x̄s) 

 end 

 x̄s = map(unthunk ∘ last, f̄_and_x̄s) # project does not support receiving InplaceableThunks 

 return NoTangent(), f̄, project(x̄s) 

 end 

 return y, sum_pullback 

 end

from chainrules.jl.

oxinabox commented on May 26, 2024

reopened as i had to revert the fix

from chainrules.jl.

Add rule for `mean(f, x; dims)` about chainrules.jl HOT 3 OPEN

Comments (3)

Related Issues (20)

Recommend Projects

React

Vue.js

Typescript

TensorFlow

Django

Laravel

D3

Recommend Topics

javascript

web

server

Machine learning

Visualization

Game

Recommend Org

Facebook

Microsoft

Google

Alibaba

D3

Tencent

	function rrule(
	config::RuleConfig{>:HasReverseMode}, ::typeof(sum), f, xs::AbstractArray; dims=:
	)
	fx_and_pullbacks = map(x->rrule_via_ad(config, f, x), xs)
	y = sum(first, fx_and_pullbacks; dims=dims)

	pullbacks = last.(fx_and_pullbacks)

	project = ProjectTo(xs)

	function sum_pullback(ȳ)
	call(f, x) = f(x)
	# if dims is :, then need only left-handed only broadcast
	broadcast_ȳ = dims isa Colon ? (ȳ,) : ȳ
	f̄_and_x̄s = call.(pullbacks, broadcast_ȳ)
	# no point thunking as most of work is in f̄_and_x̄s which we need to compute for both
	f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f
	NoTangent()
	else
	sum(first, f̄_and_x̄s)
	end
	x̄s = map(unthunk ∘ last, f̄_and_x̄s) # project does not support receiving InplaceableThunks
	return NoTangent(), f̄, project(x̄s)
	end
	return y, sum_pullback
	end