Fix heap overflow error in raw_ops.ArgMax
The API raw_ops.ArgMax implementation copies the value of arg dimension into int32_t without checking exact type which may read out of bounds.

Reported in #63070 .