Vectorize softmax and argmax nodes

[ricardoV94]

Description

The Softmax and LogSoftmax unfortunately don't support partial axes, so we have to fallback to Blockwise when more than one axis was used in the core case.

Related Issue

• Related to Vectorize follow-up #430

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[jessegrabowski]

Looks good to go, who knew vectorzing softmax would require refactoring batch axis handling. The refactor looks much cleaner and more consistent across functions.