JAX AI Stack: Port UNETR model

[jenriver]

Resolves #<issue_number_goes_here>

Reference

Checklist

• I have read the Contribution Guidelines and used pre-commit hooks to format this commit.
• I have added all the necessary unit tests for my change. (run_model.py for model usage, test_outputs.py and/or model_validation_colab.ipynb for quality).
• [] (If using an LLM) I have carefully reviewed and removed all superfluous comments or unneeded, commented-out code. Only necessary and functional code remains.
• I have signed the Contributor License Agreement (CLA).

[chapman20j]

lets consider using the ViT implementation here since just doing training and no pre-trained weights.

[chapman20j]

Can we give a more descriptive name for this layer?

[chapman20j]

I think we should also include a readme explaining what the code is doing and how it should be used.