PR #12537: [XLA:CPU][oneDNN][BugFix] Add axis support to oneDNN softmax rewriter
PR #12537: [XLA:CPU][oneDNN][BugFix] Add axis support to oneDNN softmax rewriter

Imported from GitHub PR https://github.com/openxla/xla/pull/12537

This PR fixes :
XLA Issue : https://github.com/openxla/xla/issues/11772
corresponding JAX issue : https://github.com/google/jax/issues/20856

Accuracy fix :
This PR adds support for softmax axis other than -1, previously oneDNN softmax was always executed with axis = -1 (last dimension)
Accuracy issue is observed from JAX 0.4.22 till last release, with current main branch we don't see accuracy issue as the JAX softmax HLO pattern changed (https://github.com/google/jax/pull/20643) and HLO pattern is not re-written to oneDNN softmax. Hence, this PR also adjusts the oneDNN softmax pattern to recognize new HLO pattern and rewrite into oneDNN softmax custom call.
Copybara import of the project:

--
b983bd61044bf52838661057f570e9e218803549 by Sachin Muradi <sachin.muradi@intel.com>:

Add axis support

--
07f7d14aa1f00676b21e8c1f48719fe7e56bcfd2 by Sachin Muradi <sachin.muradi@intel.com>:

address more comments + optional broadcast


Merging this change closes #12537

Reverts dbf3cd351cdac416d635a9725746f4bcd96ffd83

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/12537 from Intel-tensorflow:sachin/softmax-axis de9d6bd36272962faf76077f594649d229bb727a
