[Feature Request]: Support kFloat16 in the backward of element-wise max

[gusye1234]

Background and motivation

在如OPT, Llama等LLM中, huggingface的实现都涉及到element-wise max的操作, 且LLM大多采用float16的权重载入和amp training. 导致训练时反向传播导致错误.

API Proposal

None

API Usage

None

Alternatives

No response

Risks

No response