JackCaoG
JackCaoG
Thanks, I will try to take a look or finding someone from my team to help. nvm, this is tf2, I only knows pt/xla lol
https://github.com/pytorch/xla/blob/master/scripts/gen_lazy_tensor.py#L47 controls what got passed to the shape function. I don't think we want to pass all int64 and bool to the shape fn because in many times they don't...
actually in https://github.com/pytorch/xla/pull/3771/files I already make `bool` to be pass to all shape fns
https://github.com/pytorch/xla/pull/3771/files is merged, I think you just need to handle `dim` now
Hmm, this might be a bit tricky. The result still seems somewhat close. Not sure if the accuracy gap coming from XLA:GPU or the way we lower the `LayerNorm`
I guess the todo item here is to dump the HLO of `b` and file an issue to the XLA:GPU team.
HLO dump I get is ``` HloModule IrToHlo.44 ENTRY %IrToHlo.44 (p0.2: f32[10], p1.24: f16[3,10], p2.39: f32[10]) -> (f16[3,10]) { %constant.4 = f16[] constant(0), metadata={op_type="prim__Constant" op_name="prim__Constant" source_file="[email protected]" source_line=2501} %reshape.5 = f16[1]{0}...
Can you provide a bit more detail, like HLO snapshot you are referring to?
@ronghanghu I will give you write access 😄
Thanks for raising this issue, I will try to find someone from xla:gpu team to answer this question.