Repository: deepjavalibrary/djl
Pull request #2111 (2022-10-31 23:19:06)

Test accumulating gradient collector
This adds additional documentation and testing to define the API for the gradient collector. Specifically, it focuses on the gradient accumulating behavior and adds a restriction to creating multiple simultaneous global gradient collectors.

It also raises an issue for dealing with MXNet to accumulate gradients. I investigated both the setRequiresGradient with ADD mode and backwards with retainGrads and neither was correctly implementing the gradient accumulation behavior.
As this is not a terribly significant use case, I will just leave it as a TODO until there are any users who specifically request it.


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2111?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **71.36**% // Decreases project coverage by **`-0.72%`** :warning:
> Coverage data is based on head [(`e0499ff`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2111?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 71.65% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2111      +/-   ##
============================================
- Coverage     72.08%   71.36%   -0.73%     
- Complexity     5126     6279    +1153     
============================================
  Files           473      624     +151     
  Lines         21970    27800    +5830     
  Branches       2351     2998     +647     
============================================
+ Hits          15838    19839    +4001     
- Misses         4925     6502    +1577     
- Partials       1207     1459     +252     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2111?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...modality/cv/translator/wrapper/FileTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9GaWxlVHJhbnNsYXRvci5qYXZh) | `44.44% <ø> (ø)` | |
| [...y/cv/translator/wrapper/InputStreamTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9JbnB1dFN0cmVhbVRyYW5zbGF0b3IuamF2YQ==) | `44.44% <ø> (ø)` | |
| ... and [563 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2111/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2111?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
index 863f2160e..59b69fbe7 100644
--- a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
@@ -193,6 +193,12 @@ public abstract class BaseNDManager implements NDManager {
         throw new UnsupportedOperationException("Not supported!");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray randomPermutation(long n) {
+        throw new UnsupportedOperationException("Not supported!");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray randomUniform(float low, float high, Shape shape, DataType dataType) {
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index 72d4c08e1..26ffcd7aa 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -1176,6 +1176,14 @@ public interface NDManager extends AutoCloseable {
      */
     NDArray randomInteger(long low, long high, Shape shape, DataType dataType);
 
+    /**
+     * Returns a random permutation of integers from 0 to n - 1.
+     *
+     * @param n (int) – the upper bound (exclusive)
+     * @return a random permutation of integers from 0 to n - 1.
+     */
+    NDArray randomPermutation(long n);
+
     /**
      * Draws samples from a uniform distribution.
      *
diff --git a/api/src/main/java/ai/djl/training/GradientCollector.java b/api/src/main/java/ai/djl/training/GradientCollector.java
index edadf2158..0e16e94d8 100644
--- a/api/src/main/java/ai/djl/training/GradientCollector.java
+++ b/api/src/main/java/ai/djl/training/GradientCollector.java
@@ -21,6 +21,14 @@ import ai.djl.ndarray.NDArray;
  * performed within the try-with-resources are recorded and the variables marked. When {@link
  * #backward(NDArray) backward function} is called, gradients are collected w.r.t previously marked
  * variables.
+ *
+ * <p>The typical behavior is to open up a gradient collector during each batch and close it during
+ * the end of the batch. In this way, the gradient is reset between batches. If the gradient
+ * collector is left open for multiple calls to backwards, the gradients collected are accumulated
+ * and added together.
+ *
+ * <p>Due to limitations in most engines, the gradient collectors are global. This means that only
+ * one can be used at a time. If multiple are opened, an error will be thrown.
  */
 public interface GradientCollector extends AutoCloseable {
 
diff --git a/api/src/main/java/ai/djl/training/dataset/BatchSampler.java b/api/src/main/java/ai/djl/training/dataset/BatchSampler.java
index b44cec801..432c71905 100644
--- a/api/src/main/java/ai/djl/training/dataset/BatchSampler.java
+++ b/api/src/main/java/ai/djl/training/dataset/BatchSampler.java
@@ -73,7 +73,7 @@ public class BatchSampler implements Sampler {
 
         private long size;
         private long current;
-        private Iterator<Long> itemSampler;
+        private Iterator<Long> subSample;
 
         Iterate(RandomAccessDataset dataset) {
             current = 0;
@@ -82,7 +82,7 @@ public class BatchSampler implements Sampler {
             } else {
                 this.size = (dataset.size() + batchSize - 1) / batchSize;
             }
-            itemSampler = subSampler.sample(dataset);
+            subSample = subSampler.sample(dataset);
         }
 
         /** {@inheritDoc} */
@@ -95,8 +95,8 @@ public class BatchSampler implements Sampler {
         @Override
         public List<Long> next() {
             List<Long> batchIndices = new ArrayList<>();
-            while (itemSampler.hasNext()) {
-                batchIndices.add(itemSampler.next());
+            while (subSample.hasNext()) {
+                batchIndices.add(subSample.next());
                 if (batchIndices.size() == batchSize) {
                     break;
                 }
diff --git a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
index 1a0ac0d4e..bde8a8913 100644
--- a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
+++ b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
@@ -143,6 +143,12 @@ public final class PassthroughNDManager implements NDManager {
         throw new UnsupportedOperationException(UNSUPPORTED);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray randomPermutation(long n) {
+        throw new UnsupportedOperationException("Not supported!");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray randomUniform(float low, float high, Shape shape, DataType dataType) {
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDManager.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDManager.java
index a4bb4b2fe..7fc78b542 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDManager.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDManager.java
@@ -233,6 +233,14 @@ public class MxNDManager extends BaseNDManager {
         return invoke("_npi_random_randint", params);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray randomPermutation(long n) {
+        NDArray array = arange(0, n, 1, DataType.INT64);
+        MxOpParams params = new MxOpParams();
+        return invoke("_npi_shuffle", new NDList(array), params).singletonOrThrow();
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray randomUniform(float low, float high, Shape shape, DataType dataType) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
index e666237f0..5ef573cee 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
@@ -18,15 +18,26 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.pytorch.jni.JniUtils;
 import ai.djl.training.GradientCollector;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 /** {@code PtGradientCollector} is the PyTorch implementation of {@link GradientCollector}. */
 public final class PtGradientCollector implements GradientCollector {
 
     private boolean gradModel;
+    private static AtomicBoolean isCollecting = new AtomicBoolean();
 
     /** Constructs a new {@code PtGradientCollector} instance. */
     public PtGradientCollector() {
         gradModel = JniUtils.isGradMode();
         JniUtils.setGradMode(true);
+
+        boolean wasCollecting = isCollecting.getAndSet(true);
+        if (wasCollecting) {
+            throw new IllegalStateException(
+                    "A PtGradientCollector is already collecting. Only one can be collecting at a"
+                            + " time");
+        }
+
         zeroGradients();
     }
 
@@ -73,6 +84,7 @@ public final class PtGradientCollector implements GradientCollector {
         if (!gradModel) {
             JniUtils.setGradMode(false);
         }
+        isCollecting.set(false);
         // TODO: do some clean up if necessary
     }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
index 40af6b7a7..401da0942 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
@@ -152,6 +152,12 @@ public class PtNDManager extends BaseNDManager {
         return JniUtils.randint(this, low, high, shape, dataType, device);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray randomPermutation(long n) {
+        return JniUtils.randperm(this, n, DataType.INT64, device);
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray randomUniform(float low, float high, Shape shape, DataType dataType) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index a4bb78cb6..eae7b68af 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -1114,6 +1114,18 @@ public final class JniUtils {
                         false));
     }
 
+    public static PtNDArray randperm(
+            PtNDManager manager, long n, DataType dataType, Device device) {
+        return new PtNDArray(
+                manager,
+                PyTorchLibrary.LIB.torchRandPerm(
+                        n,
+                        dataType.ordinal(),
+                        layoutMapper(SparseFormat.DENSE, device),
+                        new int[] {PtDeviceType.toDeviceType(device), device.getDeviceId()},
+                        false));
+    }
+
     public static PtNDArray normal(
             PtNDManager manager,
             double mean,
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 378efeae4..7bab3e154 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -345,6 +345,8 @@ final class PyTorchLibrary {
             int[] device,
             boolean requiredGrad);
 
+    native long torchRandPerm(long n, int dType, int layout, int[] device, boolean requireGrad);
+
     native long torchNormal(
             double mean,
             double std,
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_random_sampling.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_random_sampling.cc
index 6a174a35b..0f53b2f1e 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_random_sampling.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_random_sampling.cc
@@ -36,6 +36,21 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchRandint(JNIE
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchRandPerm(
+    JNIEnv* env, jobject jthis, jlong jn, jint jdtype, jint jlayout, jintArray jdevice, jboolean jrequire_grad) {
+  API_BEGIN()
+  const auto options = utils::CreateTensorOptions(env, jdtype, jlayout, jdevice, jrequire_grad);
+  torch::Tensor tensor = torch::randperm(jn, options);
+  // Tensor Option for mkldnn is not working
+  // explicitly convert to mkldnn
+  if (jlayout == 2) {
+    tensor = tensor.to_mkldnn();
+  }
+  const auto* result_ptr = new torch::Tensor(tensor);
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchNormal(JNIEnv* env, jobject jthis, jdouble jmean,
     jdouble jstd, jlongArray jsizes, jint jdtype, jint jlayout, jintArray jdevice, jboolean jrequire_grad) {
   API_BEGIN()
diff --git a/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java b/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java
index acebbc93e..79c730abf 100644
--- a/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java
+++ b/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java
@@ -97,8 +97,8 @@ public final class TransferFreshFruit {
         for (Pair<String, Parameter> paramPair : baseBlock.getParameters()) {
             learningRateTrackerBuilder.put(paramPair.getValue().getId(), 0.1f * lr);
         }
-        Optimizer optimizer =
-                Adam.builder().optLearningRateTracker(learningRateTrackerBuilder.build()).build();
+        FixedPerVarTracker learningRateTracker = learningRateTrackerBuilder.build();
+        Optimizer optimizer = Adam.builder().optLearningRateTracker(learningRateTracker).build();
         config.optOptimizer(optimizer);
 
         Trainer trainer = model.newTrainer(config);
@@ -131,7 +131,7 @@ public final class TransferFreshFruit {
         float[] mean = {0.485f, 0.456f, 0.406f};
         float[] std = {0.229f, 0.224f, 0.225f};
 
-        // If the user wants to use local repository, then the dataset can be loaded as follows
+        // If users want to use local repository, then the dataset can be loaded as follows
         // Repository repository = Repository.newInstance("banana", Paths.get(LOCAL_FOLDER/{train OR
         // test}));
         // FruitsFreshAndRotten dataset =


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
index 89a2dadcc..403ca2c50 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
@@ -24,6 +24,7 @@ import java.io.IOException;
 public class TrainAirfoilWithTabNetTest {
     @Test
     public void testTrainAirfoilWithTabNet() throws TranslateException, IOException {
+        TestRequirements.nightly();
         TestRequirements.engine("MXNet", "PyTorch");
         String[] args = new String[] {"-g", "1", "-e", "20", "-b", "32"};
         TrainingResult result = TrainAirfoilWithTabNet.runExample(args);
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java
index 37e18db1b..00b346eab 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java
@@ -394,6 +394,21 @@ public class NDArrayCreationOpTest {
         }
     }
 
+    @Test
+    public void testRandomPermutation() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            long size = 3;
+            NDArray array = manager.randomPermutation(size);
+            Assert.assertTrue(
+                    array.contentEquals(manager.create(new long[] {0, 1, 2}))
+                            || array.contentEquals(manager.create(new long[] {0, 2, 1}))
+                            || array.contentEquals(manager.create(new long[] {1, 0, 2}))
+                            || array.contentEquals(manager.create(new long[] {1, 2, 0}))
+                            || array.contentEquals(manager.create(new long[] {2, 0, 1}))
+                            || array.contentEquals(manager.create(new long[] {2, 1, 0})));
+        }
+    }
+
     @Test
     public void testRandomUniform() {
         try (NDManager manager = NDManager.newBaseManager()) {
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java b/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
index 1baf9dd37..92f491ba5 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
@@ -121,6 +121,46 @@ public class GradientCollectorIntegrationTest {
         }
     }
 
+    /** Tests that the gradients do accumulate within the same gradient collector. */
+    @Test
+    public void testAccumulateGradients() {
+        // TODO: MXNet support for accumulating gradients does not currently work
+        TestRequirements.notEngine("MXNet");
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray a = manager.create(0.0f);
+            a.setRequiresGradient(true);
+
+            try (GradientCollector gc = Engine.getInstance().newGradientCollector()) {
+                for (int i = 1; i <= 3; i++) {
+                    NDArray b = a.mul(2);
+                    gc.backward(b);
+                    Assert.assertEquals(a.getGradient().getFloat(), 2.0f * i);
+                }
+            }
+        }
+    }
+
+    /**
+     * Ensures that a gradient collector does not start when one is already created because they are
+     * global.
+     */
+    @Test
+    @SuppressWarnings({"try", "PMD.UseTryWithResources"})
+    public void testMultipleGradientCollectors() {
+        Assert.assertThrows(
+                () -> {
+                    GradientCollector gc2 = null;
+                    try (GradientCollector gc = Engine.getInstance().newGradientCollector()) {
+                        gc2 = Engine.getInstance().newGradientCollector();
+                        gc2.close();
+                    } finally {
+                        if (gc2 != null) {
+                            gc2.close();
+                        }
+                    }
+                });
+    }
+
     @Test
     public void testFreezeParameters() {
         try (Model model = Model.newInstance("model")) {