[torchlib] Fix _log_softmax

[justinchuby]

Fix _log_softmax by moving the IsScalar call to the top so it can be eagerly evaluated.

Also specify the squeeze axis explicitly to improve compatibility with ORT: microsoft/onnxruntime#21661

This should fix a runtime error in XGLMForCausalLM

[codecov]

Codecov Report

Attention: Patch coverage is 37.50000% with 5 lines in your changes missing coverage. Please review.

Project coverage is 75.02%. Comparing base (cf5ddd9) to head (daeb3b3).
Report is 1 commits behind head on main.

Files	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	37.50%	2 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1789      +/-   ##
==========================================
- Coverage   75.04%   75.02%   -0.03%     
==========================================
  Files         245      245              
  Lines       26512    26518       +6     
  Branches     4834     4836       +2     
==========================================
- Hits        19897    19895       -2     
- Misses       5683     5688       +5     
- Partials      932      935       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[github-actions]

Test Results

     24 files  ±     0       24 suites  ±0   3h 26m 17s ⏱️ + 1m 47s
 13 506 tests  -  2 183   11 936 ✅  - 1 746    1 543 💤  -    434   27 ❌  - 3 
481 900 runs  +18 070  100 122 ✅ +2 632  381 550 💤 +15 441  228 ❌  - 3 

For more details on these failures, see this check.

Results for commit daeb3b3. ± Comparison against base commit cf5ddd9.

This pull request removes 2183 tests.

onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_basic_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_doc_string
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_if_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_if_loop_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_loop_defs
onnxscript._internal.analysis_test.TestExposedUses ‑ test_basic
onnxscript._internal.analysis_test.TestExposedUses ‑ test_called_function
onnxscript._internal.analysis_test.TestExposedUses ‑ test_doc_string
onnxscript._internal.analysis_test.TestExposedUses ‑ test_for_loop
onnxscript._internal.analysis_test.TestExposedUses ‑ test_if
…

♻️ This comment has been updated with latest results.