Add position encoding for AxialAttention

[jacobbieker]

Pull Request

Description

Please delete the italicised instruction text! Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes #

How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

• [ ] Yes

If your changes affect data processing, have you plotted any changes? i.e. have you done a quick sanity check?

• [ ] Yes

Checklist:

• [ ] My code follows OCF's coding style guidelines
• [ ] I have performed a self-review of my own code
• [ ] I have made corresponding changes to the documentation
• [ ] I have added tests that prove my fix is effective or that my feature works
• [ ] I have checked my code and corrected any misspellings