[Android] Move training part to its own package

Author: kirklandsign

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/TrainingModuleE2ETest.kt

@@ -20,6 +20,8 @@ import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.pytorch.executorch.TestFileUtils.getTestFilePath
+import org.pytorch.executorch.training.SGD
+import org.pytorch.executorch.training.TrainingModule
 import kotlin.random.Random
 import kotlin.test.assertContains
 

extension/android/executorch_android/src/main/java/org/pytorch/executorch/SGD.java renamed to extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/SGD.java

@@ -1,18 +1,22 @@
 /*
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- * All rights reserved.
  *
- * This source code is licensed under the BSD-style license found in the
- * LICENSE file in the root directory of this source tree.
+ *  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *  * All rights reserved.
+ *  *
+ *  * This source code is licensed under the BSD-style license found in the
+ *  * LICENSE file in the root directory of this source tree.
+ *
+ *
  */
 
-package org.pytorch.executorch;
+package org.pytorch.executorch.training;
 
 import com.facebook.jni.HybridData;
 import com.facebook.jni.annotations.DoNotStrip;
 import com.facebook.soloader.nativeloader.NativeLoader;
 import com.facebook.soloader.nativeloader.SystemDelegate;
 import java.util.Map;
+import org.pytorch.executorch.Tensor;
 import org.pytorch.executorch.annotations.Experimental;
 
 /**
@@ -62,7 +66,7 @@ private SGD(
    * @param dampening The dampening value
    * @param weightDecay The weight decay value
    * @param nesterov Whether to use Nesterov momentum
-   * @return new {@link org.pytorch.executorch.SGD} object
+   * @return new {@link SGD} object
    */
   public static SGD create(
       Map<String, Tensor> namedParameters,
@@ -79,7 +83,7 @@ public static SGD create(
    *
    * @param namedParameters Map of parameter names to tensors to be optimized
    * @param learningRate The learning rate for the optimizer
-   * @return new {@link org.pytorch.executorch.SGD} object
+   * @return new {@link SGD} object
    */
   public static SGD create(Map<String, Tensor> namedParameters, double learningRate) {
     return create(namedParameters, learningRate, 0.0, 0.0, 0.0, false);

extension/android/executorch_android/src/main/java/org/pytorch/executorch/TrainingModule.java renamed to extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/TrainingModule.java

@@ -1,12 +1,15 @@
 /*
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- * All rights reserved.
  *
- * This source code is licensed under the BSD-style license found in the
- * LICENSE file in the root directory of this source tree.
+ *  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *  * All rights reserved.
+ *  *
+ *  * This source code is licensed under the BSD-style license found in the
+ *  * LICENSE file in the root directory of this source tree.
+ *
+ *
  */
 
-package org.pytorch.executorch;
+package org.pytorch.executorch.training;
 
 import android.util.Log;
 import com.facebook.jni.HybridData;
@@ -16,6 +19,8 @@
 import java.io.File;
 import java.util.HashMap;
 import java.util.Map;
+import org.pytorch.executorch.EValue;
+import org.pytorch.executorch.Tensor;
 import org.pytorch.executorch.annotations.Experimental;
 
 /**
@@ -48,7 +53,7 @@ private TrainingModule(String moduleAbsolutePath, String dataAbsolutePath) {
    *
    * @param modelPath path to file that contains the serialized ExecuTorch module.
    * @param dataPath path to file that contains the ExecuTorch module external weights.
-   * @return new {@link org.pytorch.executorch.TrainingModule} object which owns the model module.
+   * @return new {@link TrainingModule} object which owns the model module.
    */
   public static TrainingModule load(final String modelPath, final String dataPath) {
     File modelFile = new File(modelPath);
@@ -67,7 +72,7 @@ public static TrainingModule load(final String modelPath, final String dataPath)
    *
    * @param modelPath path to file that contains the serialized ExecuTorch module. This PTE does not
    *     rely on external weights.
-   * @return new {@link org.pytorch.executorch.TrainingModule} object which owns the model module.
+   * @return new {@link TrainingModule} object which owns the model module.
    */
   public static TrainingModule load(final String modelPath) {
     File modelFile = new File(modelPath);

extension/android/jni/jni_layer_training.cpp

@@ -67,7 +67,7 @@ class ExecuTorchTrainingJni
 
  public:
   constexpr static auto kJavaDescriptor =
-      "Lorg/pytorch/executorch/TrainingModule;";
+      "Lorg/pytorch/executorch/training/TrainingModule;";
 
   ExecuTorchTrainingJni(
       facebook::jni::alias_ref<jstring> modelPath,
@@ -226,7 +226,8 @@ class ExecuTorchTrainingJni
 
 class SGDHybrid : public facebook::jni::HybridClass<SGDHybrid> {
  public:
-  constexpr static const char* kJavaDescriptor = "Lorg/pytorch/executorch/SGD;";
+  constexpr static const char* kJavaDescriptor =
+      "Lorg/pytorch/executorch/training/SGD;";
 
   static facebook::jni::local_ref<jhybriddata> initHybrid(
       facebook::jni::alias_ref<jclass>,