Skip to content

Instantly share code, notes, and snippets.

@apivovarov
Created November 7, 2023 00:57
Show Gist options
  • Save apivovarov/b6123d3e99d26cc7507b8b04f82b0781 to your computer and use it in GitHub Desktop.
Save apivovarov/b6123d3e99d26cc7507b8b04f82b0781 to your computer and use it in GitHub Desktop.
HloModule xla_computation_ff, entry_computation_layout={(f32[4,1000]{1,0})->(f32[4,1000]{1,0})}
region_0.4 {
Arg_0.5 = f32[] parameter(0)
Arg_1.6 = f32[] parameter(1)
ROOT maximum.7 = f32[] maximum(Arg_0.5, Arg_1.6)
}
region_1.15 {
Arg_0.16 = f32[] parameter(0)
Arg_1.17 = f32[] parameter(1)
ROOT add.18 = f32[] add(Arg_0.16, Arg_1.17)
}
ENTRY main.26 {
Arg_0.1 = f32[4,1000]{1,0} parameter(0)
constant.3 = f32[] constant(-inf)
reduce.8 = f32[4]{0} reduce(Arg_0.1, constant.3), dimensions={1}, to_apply=region_0.4
reshape.9 = f32[4,1]{1,0} reshape(reduce.8)
broadcast.10 = f32[4,1]{1,0} broadcast(reshape.9), dimensions={0,1}
reshape.11 = f32[4]{0} reshape(broadcast.10)
broadcast.12 = f32[4,1000]{1,0} broadcast(reshape.11), dimensions={0}
subtract.13 = f32[4,1000]{1,0} subtract(Arg_0.1, broadcast.12)
exponential.14 = f32[4,1000]{1,0} exponential(subtract.13)
constant.2 = f32[] constant(0)
reduce.19 = f32[4]{0} reduce(exponential.14, constant.2), dimensions={1}, to_apply=region_1.15
reshape.20 = f32[4,1]{1,0} reshape(reduce.19)
broadcast.21 = f32[4,1]{1,0} broadcast(reshape.20), dimensions={0,1}
reshape.22 = f32[4]{0} reshape(broadcast.21)
broadcast.23 = f32[4,1000]{1,0} broadcast(reshape.22), dimensions={0}
divide.24 = f32[4,1000]{1,0} divide(exponential.14, broadcast.23)
ROOT tuple.25 = (f32[4,1000]{1,0}) tuple(divide.24)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment