[RFC] Add TrainingModule and SGD JNI

[georgehong]

As title, adds wrappers together with unit test based on XOR train.cpp example.

Summary

Adds JNI for SGD and TrainingModule, including a unit test that mirrors train.cpp for a simple XOR example. Also makes the following change:

• Refactor jni_layer.cpp JTensor <--> Tensor conversion to be a general TensorHybrid utility. This is useful for TrainingModule classes that move maps of Tensors around.

Training dependencies are already enabled for Java JNI library, so we skip adding additional guard flags.

Test plan

Followed steps in README for extension/android, namely:

sh scripts/build_android_library.sh
sh executorch_android/android_test_setup.sh // Now creates xor.ptd and xor.pte test dependencies

./gradlew :executorch_android:connectedAndroidTest // Added unit test to check toy model convergence loss < 0.01

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/12188

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Cancelled Job, 1 Unrelated Failure

As of commit 19da644 with merge base 9905026 ():

NEW FAILURE - The following job has failed:

• pull / android / build-llm-demo / linux-job (gh)
  RuntimeError: Command docker exec -t d7af12d21be66c45a2dc7f921a22fc66b0250747f523aff05f1fd9c45d3140ea /exec failed with exit code 139

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-models-linux (mobilebert, xnnpack-quantization-delegation, linux.2xlarge) / linux-job (gh)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-models-linux (mobilebert, portable, linux.2xlarge) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[facebook-github-bot]

@georgehong has imported this pull request. If you are a Meta employee, you can view this in D77702768.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[facebook-github-bot]

@georgehong has imported this pull request. If you are a Meta employee, you can view this in D77702768.

[facebook-github-bot]

@georgehong has imported this pull request. If you are a Meta employee, you can view this in D77702768.