allow customized head_dim

[helunwencser]

Stack from ghstack (oldest at bottom):

• -> #6872

This is for resolving the ask in this post.

Similar change in HF: https://github.com/huggingface/transformers/pull/32502

Differential Revision: D65974454

[pytorch-bot[bot]]

:link: Helpful Links

:test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/6872

• :page_facing_up: Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

:white_check_mark: No Failures

As of commit 811ad8bc46e5f792f41a9e1ab83de7a42354981d with merge base fbcc9a13f875113ab2a0e0f8a67a6586f1fd3bda (): :green_heart: Looks good so far! There are no failures yet. :green_heart:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D65974454

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D65974454

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D65974454

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D65974454

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D65974454

[helunwencser]

@pytorchbot label "topic: not user facing"