Fix UNet segmentation example

[xrhd]

Resolves #113 #73

Reference

• jax.extend.backend.get_backend
• Orbax checkpoints

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]

Hi @xrhd. Thanks for this commit. Just leaving some brief comments.

1. Could you remove emojis from the notebook
2. Could you update the intro information to reflect that this is no longer a work in progress
3. In section 4, the script re-loads the weights and then trains again. The further trained weights aren't used in the final inference outputs. Could you just use the weights from section 4 in section 6?

[jenriver]

Also, please update the branch to make sure latest codes are in. Thanks!

[jenriver]

Hi @xrhd, thanks for the update!

It looks like the pre-commit check is failing because we recently added jupytext to the repository. We use this to automatically sync .ipynb files with a paired .md file, which makes reviewing notebook diffs much easier.
Since your environment might not have the updated dependencies yet, the .md file wasn't generated when you saved the notebook. Could you please run the following steps to fix the sync?

# 1. Update your dev dependencies (to install jupytext)
pip install -e ".[dev]"

# 2. Run pre-commit manually
pre-commit run --all-files

This will likely modify the .ipynb metadata slightly and generate/update the corresponding .md file. Please add those changes to your commit, and the checks should pass.

[xrhd]

Hi @chapman20j and @jenriver,

Thanks for the review and feedback!

First of all, sorry for the delay in getting back to this. I've gone ahead and applied the changes you requested:

• Updated the intro to reflect that it's no longer a WIP.
• Adjusted the logic in sections 4 and 6 to ensure the correct weights are used for inference.

I also followed the steps to fix the pre-commit and jupytext sync, so the checks should be passing now.

[chapman20j]

Hi. Sorry for the delays in getting this accepted. I took a look and it looks good to me. Could you squash this commit into 1 commit?

[xrhd]

Hi. Sorry for the delays in getting this accepted. I took a look and it looks good to me. Could you squash this commit into 1 commit?

Squashed and rebased @chapman20j :)