[Semi-Auto] Add softmax infer_backward rule #56502
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
|
Sorry to inform you that 05ef040's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually. |
pkuzyc
force-pushed
the
softmax_backward
branch
from
05ef040 to
57c43f7
Compare
| // sharding on softmax_axis is not supported now, so set | ||
| // the notation on softmax_axis to '1' so that we can set | ||
| // its dim mapping to -1 | ||
| x_axes[axis] = '1'; |
There was a problem hiding this comment.
change the einsum_string to "1" to ensure the unsharded maybe not a good idea here, since it would disconnect the relation of softmax_axis between input and output when print the vlog for debug:
"ab1" --> "abc", later developer might not realize that axis "1" and "c" are the same axis in two tensors.
it would be better that:
first: out_dims_mapping[axis] = -1;
then: keep the input_axes unchanged.
PR types
Function optimization
PR changes
Others
Description
Pcard-70448
Add infer backward SPMD rule for softmax. It infers inputs' dims mappings as following:
For non-normalized axes, set their dims mappings equal to the outputs'. For the normalized axis, set it to -1.