Allow softmax and log_softmax to operate on any dimension (#5694)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/5694

Use the portable implementation for `dim != -1`, simplified to only floats.

Make the optimized cases fallback to the CPU impl.

Reviewed By: manuelcandales

Differential Revision: D63430960

fbshipit-source-id: ec2d56a60d2b064744383a2f1266dee7600fa884