refactor: use xlogx and xlogy from LogExpFunctions

[avik-pal]

can't go forward with this, these functions are not safe for broadcasting on GPUs https://buildkite.com/julialang/lux-dot-jl/builds/3368#0190f65f-7baa-463e-8072-5c3fb2fd0b0b/152-307

[github-actions]

Benchmark Results

Benchmark suite	Current: 3b89ccf	Previous: f3159bd	Ratio
Dense(2 => 2)/cpu/reverse/ReverseDiff (compiled)/(2, 128)	3734.5 ns	4340.625 ns	0.86
Dense(2 => 2)/cpu/reverse/Zygote/(2, 128)	7358.833333333333 ns	8991.833333333334 ns	0.82
Dense(2 => 2)/cpu/reverse/Tracker/(2, 128)	21220 ns	25788 ns	0.82
Dense(2 => 2)/cpu/reverse/ReverseDiff/(2, 128)	9826.4 ns	12129.333333333334 ns	0.81
Dense(2 => 2)/cpu/reverse/Flux/(2, 128)	9087 ns	10748.1 ns	0.85
Dense(2 => 2)/cpu/reverse/SimpleChains/(2, 128)	4485.875 ns	4954.25 ns	0.91
Dense(2 => 2)/cpu/reverse/Enzyme/(2, 128)	4676.25 ns	5425.125 ns	0.86
Dense(2 => 2)/cpu/forward/NamedTuple/(2, 128)	1114.1572327044025 ns	1359.1832061068703 ns	0.82
Dense(2 => 2)/cpu/forward/ComponentArray/(2, 128)	1186.9925373134329 ns	1401.2664233576643 ns	0.85
Dense(2 => 2)/cpu/forward/Flux/(2, 128)	1770.826923076923 ns	2176.5833333333335 ns	0.81
Dense(2 => 2)/cpu/forward/SimpleChains/(2, 128)	180.3095238095238 ns	179.46778711484595 ns	1.00
Dense(20 => 20)/cpu/reverse/ReverseDiff (compiled)/(20, 128)	17322 ns	21761 ns	0.80
Dense(20 => 20)/cpu/reverse/Zygote/(20, 128)	16872 ns	20408 ns	0.83
Dense(20 => 20)/cpu/reverse/Tracker/(20, 128)	37280 ns	42780 ns	0.87
Dense(20 => 20)/cpu/reverse/ReverseDiff/(20, 128)	28964 ns	32301 ns	0.90
Dense(20 => 20)/cpu/reverse/Flux/(20, 128)	21540 ns	22652 ns	0.95
Dense(20 => 20)/cpu/reverse/SimpleChains/(20, 128)	17392.5 ns	18654 ns	0.93
Dense(20 => 20)/cpu/reverse/Enzyme/(20, 128)	25538 ns	28283 ns	0.90
Dense(20 => 20)/cpu/forward/NamedTuple/(20, 128)	3888.625 ns	4363.8125 ns	0.89
Dense(20 => 20)/cpu/forward/ComponentArray/(20, 128)	3961.125 ns	4459.5 ns	0.89
Dense(20 => 20)/cpu/forward/Flux/(20, 128)	4872 ns	5581.928571428572 ns	0.87
Dense(20 => 20)/cpu/forward/SimpleChains/(20, 128)	1655.1 ns	1850.5 ns	0.89
Conv((3, 3), 3 => 3)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 3, 128)	40760948 ns	45154820.5 ns	0.90
Conv((3, 3), 3 => 3)/cpu/reverse/Zygote/(64, 64, 3, 128)	58630590.5 ns	62378075.5 ns	0.94
Conv((3, 3), 3 => 3)/cpu/reverse/Tracker/(64, 64, 3, 128)	76151396 ns	82237576.5 ns	0.93
Conv((3, 3), 3 => 3)/cpu/reverse/ReverseDiff/(64, 64, 3, 128)	81998162 ns	92405677.5 ns	0.89
Conv((3, 3), 3 => 3)/cpu/reverse/Flux/(64, 64, 3, 128)	75223403.5 ns	78467695 ns	0.96
Conv((3, 3), 3 => 3)/cpu/reverse/SimpleChains/(64, 64, 3, 128)	12029690 ns	11963195.5 ns	1.01
Conv((3, 3), 3 => 3)/cpu/reverse/Enzyme/(64, 64, 3, 128)	91402818 ns	91650799 ns	1.00
Conv((3, 3), 3 => 3)/cpu/forward/NamedTuple/(64, 64, 3, 128)	7669704 ns	7710446.5 ns	0.99
Conv((3, 3), 3 => 3)/cpu/forward/ComponentArray/(64, 64, 3, 128)	7605872 ns	7584631 ns	1.00
Conv((3, 3), 3 => 3)/cpu/forward/Flux/(64, 64, 3, 128)	11945823 ns	12087804.5 ns	0.99
Conv((3, 3), 3 => 3)/cpu/forward/SimpleChains/(64, 64, 3, 128)	6391577.5 ns	6389570.5 ns	1.00
vgg16/cpu/reverse/Zygote/(32, 32, 3, 16)	706066648.5 ns	691470864.5 ns	1.02
vgg16/cpu/reverse/Zygote/(32, 32, 3, 64)	2594444144 ns	2612057812 ns	0.99
vgg16/cpu/reverse/Zygote/(32, 32, 3, 2)	135741873 ns	151710246 ns	0.89
vgg16/cpu/reverse/Tracker/(32, 32, 3, 16)	864869891 ns	847276341 ns	1.02
vgg16/cpu/reverse/Tracker/(32, 32, 3, 64)	3200921122 ns	2998224779 ns	1.07
vgg16/cpu/reverse/Tracker/(32, 32, 3, 2)	240317951.5 ns	231162547 ns	1.04
vgg16/cpu/reverse/Flux/(32, 32, 3, 16)	816953875 ns	712811787.5 ns	1.15
vgg16/cpu/reverse/Flux/(32, 32, 3, 64)	2545771395 ns	2949177680 ns	0.86
vgg16/cpu/reverse/Flux/(32, 32, 3, 2)	127485771 ns	139771741 ns	0.91
vgg16/cpu/forward/NamedTuple/(32, 32, 3, 16)	174430006.5 ns	175403229 ns	0.99
vgg16/cpu/forward/NamedTuple/(32, 32, 3, 64)	649400570.5 ns	643986392.5 ns	1.01
vgg16/cpu/forward/NamedTuple/(32, 32, 3, 2)	34520164 ns	45299778 ns	0.76
vgg16/cpu/forward/ComponentArray/(32, 32, 3, 16)	164755022 ns	165572824 ns	1.00
vgg16/cpu/forward/ComponentArray/(32, 32, 3, 64)	658744577 ns	648012066 ns	1.02
vgg16/cpu/forward/ComponentArray/(32, 32, 3, 2)	40083686 ns	35881786 ns	1.12
vgg16/cpu/forward/Flux/(32, 32, 3, 16)	218808171 ns	228284686 ns	0.96
vgg16/cpu/forward/Flux/(32, 32, 3, 64)	762249988.5 ns	824337722 ns	0.92
vgg16/cpu/forward/Flux/(32, 32, 3, 2)	36076727 ns	36049946 ns	1.00
Conv((3, 3), 64 => 64)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 64, 128)	1245208126 ns	1305860557.5 ns	0.95
Conv((3, 3), 64 => 64)/cpu/reverse/Zygote/(64, 64, 64, 128)	1883004023 ns	1872497449 ns	1.01
Conv((3, 3), 64 => 64)/cpu/reverse/Tracker/(64, 64, 64, 128)	2365936471 ns	2342556041 ns	1.01
Conv((3, 3), 64 => 64)/cpu/reverse/ReverseDiff/(64, 64, 64, 128)	2416275974 ns	2488038610 ns	0.97
Conv((3, 3), 64 => 64)/cpu/reverse/Flux/(64, 64, 64, 128)	1947501810 ns	1951631856 ns	1.00
Conv((3, 3), 64 => 64)/cpu/reverse/Enzyme/(64, 64, 64, 128)	2181618043 ns	2087636119 ns	1.05
Conv((3, 3), 64 => 64)/cpu/forward/NamedTuple/(64, 64, 64, 128)	340900455 ns	338033674 ns	1.01
Conv((3, 3), 64 => 64)/cpu/forward/ComponentArray/(64, 64, 64, 128)	333037472 ns	332696625 ns	1.00
Conv((3, 3), 64 => 64)/cpu/forward/Flux/(64, 64, 64, 128)	468886350 ns	359960422 ns	1.30
Conv((3, 3), 1 => 1)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 1, 128)	11823237 ns	11686611 ns	1.01
Conv((3, 3), 1 => 1)/cpu/reverse/Zygote/(64, 64, 1, 128)	18142588.5 ns	17986098.5 ns	1.01
Conv((3, 3), 1 => 1)/cpu/reverse/Tracker/(64, 64, 1, 128)	19222015 ns	19052140 ns	1.01
Conv((3, 3), 1 => 1)/cpu/reverse/ReverseDiff/(64, 64, 1, 128)	23924727 ns	23778464 ns	1.01
Conv((3, 3), 1 => 1)/cpu/reverse/Flux/(64, 64, 1, 128)	17915780.5 ns	17775034.5 ns	1.01
Conv((3, 3), 1 => 1)/cpu/reverse/SimpleChains/(64, 64, 1, 128)	1162320.5 ns	1157462 ns	1.00
Conv((3, 3), 1 => 1)/cpu/reverse/Enzyme/(64, 64, 1, 128)	22991391 ns	22924969.5 ns	1.00
Conv((3, 3), 1 => 1)/cpu/forward/NamedTuple/(64, 64, 1, 128)	2282387 ns	2406465 ns	0.95
Conv((3, 3), 1 => 1)/cpu/forward/ComponentArray/(64, 64, 1, 128)	2209775 ns	2242648.5 ns	0.99
Conv((3, 3), 1 => 1)/cpu/forward/Flux/(64, 64, 1, 128)	2065866 ns	2065648 ns	1.00
Conv((3, 3), 1 => 1)/cpu/forward/SimpleChains/(64, 64, 1, 128)	205565 ns	197910 ns	1.04
Dense(200 => 200)/cpu/reverse/ReverseDiff (compiled)/(200, 128)	294531 ns	290874 ns	1.01
Dense(200 => 200)/cpu/reverse/Zygote/(200, 128)	265336 ns	264424 ns	1.00
Dense(200 => 200)/cpu/reverse/Tracker/(200, 128)	366676 ns	364661 ns	1.01
Dense(200 => 200)/cpu/reverse/ReverseDiff/(200, 128)	408735 ns	405888 ns	1.01
Dense(200 => 200)/cpu/reverse/Flux/(200, 128)	274614 ns	272550 ns	1.01
Dense(200 => 200)/cpu/reverse/SimpleChains/(200, 128)	406561 ns	408042 ns	1.00
Dense(200 => 200)/cpu/reverse/Enzyme/(200, 128)	396943 ns	394657 ns	1.01
Dense(200 => 200)/cpu/forward/NamedTuple/(200, 128)	81443 ns	80941 ns	1.01
Dense(200 => 200)/cpu/forward/ComponentArray/(200, 128)	81913 ns	81403 ns	1.01
Dense(200 => 200)/cpu/forward/Flux/(200, 128)	86692 ns	86461 ns	1.00
Dense(200 => 200)/cpu/forward/SimpleChains/(200, 128)	104446 ns	104525 ns	1.00
Conv((3, 3), 16 => 16)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 16, 128)	191157666 ns	206243721 ns	0.93
Conv((3, 3), 16 => 16)/cpu/reverse/Zygote/(64, 64, 16, 128)	330571488.5 ns	327524051.5 ns	1.01
Conv((3, 3), 16 => 16)/cpu/reverse/Tracker/(64, 64, 16, 128)	404128258 ns	410738774 ns	0.98
Conv((3, 3), 16 => 16)/cpu/reverse/ReverseDiff/(64, 64, 16, 128)	461916847.5 ns	439152524.5 ns	1.05
Conv((3, 3), 16 => 16)/cpu/reverse/Flux/(64, 64, 16, 128)	387077386 ns	382717413 ns	1.01
Conv((3, 3), 16 => 16)/cpu/reverse/SimpleChains/(64, 64, 16, 128)	330192292 ns	325870578.5 ns	1.01
Conv((3, 3), 16 => 16)/cpu/reverse/Enzyme/(64, 64, 16, 128)	480516350 ns	455578258 ns	1.05
Conv((3, 3), 16 => 16)/cpu/forward/NamedTuple/(64, 64, 16, 128)	47405090 ns	47303828.5 ns	1.00
Conv((3, 3), 16 => 16)/cpu/forward/ComponentArray/(64, 64, 16, 128)	46728378 ns	46793042.5 ns	1.00
Conv((3, 3), 16 => 16)/cpu/forward/Flux/(64, 64, 16, 128)	51294940 ns	57476949 ns	0.89
Conv((3, 3), 16 => 16)/cpu/forward/SimpleChains/(64, 64, 16, 128)	27829814 ns	27955898 ns	1.00
Dense(2000 => 2000)/cpu/reverse/ReverseDiff (compiled)/(2000, 128)	19022622 ns	18906265 ns	1.01
Dense(2000 => 2000)/cpu/reverse/Zygote/(2000, 128)	19601144 ns	19556819 ns	1.00
Dense(2000 => 2000)/cpu/reverse/Tracker/(2000, 128)	23445038 ns	23268022 ns	1.01
Dense(2000 => 2000)/cpu/reverse/ReverseDiff/(2000, 128)	24162340 ns	24118160 ns	1.00
Dense(2000 => 2000)/cpu/reverse/Flux/(2000, 128)	19671987 ns	19557901 ns	1.01
Dense(2000 => 2000)/cpu/reverse/Enzyme/(2000, 128)	20926796 ns	20840171 ns	1.00
Dense(2000 => 2000)/cpu/forward/NamedTuple/(2000, 128)	6525363 ns	6539910.5 ns	1.00
Dense(2000 => 2000)/cpu/forward/ComponentArray/(2000, 128)	6511266 ns	6488719 ns	1.00
Dense(2000 => 2000)/cpu/forward/Flux/(2000, 128)	6502389.5 ns	6511897 ns	1.00

This comment was automatically generated by workflow using github-action-benchmark.