Repository: deepjavalibrary/djl
Pull request #2873 (2023-11-25 10:53:43)

Minor fixes to improve Apple Silicon MPS support
## Description ##

This PR contains minor changes to `NDArrayEx` and `Classifications` to avoid converting float32 to float64 when not required.

I have been using this workarounds [here](https://github.com/qupath/qupath-extension-wsinfer/blob/main/src/main/java/qupath/ext/wsinfer/MpsSupport.java) to enable MPS support on Apple Silicon for image classification tasks using PyTorch, and have seen a major speedup (as expected).

I added tests as a separate commit in case they are useful, but the relevant changes will only be noticed when running on Apple Silicon.

> Update: Tests now moved out of separate files and into `MpsTest`

This PR relates to https://github.com/deepjavalibrary/djl/pull/2044

Happy to add any required copyright statements to the tests (or remove them entirely), or look for other MPS-related issues if this is useful. I couldn't get the full tests to run on Apple Silicon as it currently stands due to TensorFlow not being supported.

petebankhead: I have added a commit to address https://github.com/deepjavalibrary/djl/issues/2504

When combined with the other changes, it is possible to train using Mnist and FashionMnist with MPS as the device.

There seems to be a deeper issue though, related to the use of `NDArray.set(buf)`. This appears to be the source of the reported error. The error occurred only when `NDarray.toType(DataType.FLOAT32, false)` was called, but `toType` appears to work if `set(buf)` is avoided. So the workaround here is to use the buffer at the time the array is created, rather than afterwards.

Here is a quick test to illustrate what happens:
```java
    private static void quickTest() {
        Device device = Device.of("mps", -1);
        try (NDManager manager = NDManager.newBaseManager(device)) {
            // This works on MPS
            try (NDArray array = manager.create(ByteBuffer.wrap(new byte[]{127}), new Shape(1, 1, 1, 1), DataType.UINT8)) {
                NDArray output = array.toType(DataType.FLOAT32, false);
                System.out.println(output);
            }
            // This fails on MPS (SIGSEGV)
            try (NDArray array = manager.create(new Shape(1, 1, 1, 1), DataType.UINT8)) {
                array.set(new byte[]{127});
                NDArray output = array.toType(DataType.FLOAT32, false);
                System.out.println(output);
            }
        }
    }
```

To train, I used the following code based on the [MNIST tutorial](https://github.com/deepjavalibrary/djl/blob/5ab2685b01b13ccceb99a0ca1667112b2b68e4f3/jupyter/tutorial/02_train_your_first_model.ipynb):
```java
        Device device = Device.of("mps", 0);
        try (NDManager manager = NDManager.newBaseManager(device)) {
            Mnist mnist = Mnist.builder()
                    .optDevice(device)
                    .optManager(manager)
                    .setSampling(batchSize, true)
                    .build();

            mnist.prepare();

            Model model = Model.newInstance("mlp", device);
            model.setBlock(new Mlp(28 * 28, 10, new int[] {128, 64}));

            DefaultTrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
                    .optDevices(new Device[]{device})
                    .addEvaluator(new Accuracy()) // Use accuracy so we humans can understand how accurate the model is
                    .addTrainingListeners(TrainingListener.Defaults.logging());

            // Now that we have our training configuration, we should create a new trainer for our model
            Trainer trainer = model.newTrainer(config);
            System.out.println("DEVICES: " + Arrays.toString(trainer.getDevices()));

            trainer.setMetrics(new Metrics());
            trainer.initialize(new Shape(1, 28*28));

            // Deep learning is typically trained in epochs where each epoch trains the model on each item in the dataset once.
            int epoch = 20;

            long start = System.currentTimeMillis();
            RandomAccessDataset split[] = mnist.randomSplit(6, 4);
            EasyTrain.fit(trainer, epoch, split[0], split[1]);
            long end = System.currentTimeMillis();

            System.out.println("Time: " + (end - start) / 1000.0);
        } catch (Exception e) {
            e.printStackTrace();
        }
```

Using `Device.of("mps", 0)` was important; if I change to `Device.of("mps", -1)` I see

```
java.lang.NullPointerException: Cannot invoke "java.lang.Integer.intValue()" because the return value of "java.util.Map.get(Object)" is null
	at ai.djl.training.ParameterStore.getValue(ParameterStore.java:104)
	at ai.djl.nn.core.Linear.forwardInternal(Linear.java:93)
	at ai.djl.nn.AbstractBaseBlock.forwardInternal(AbstractBaseBlock.java:128)
	at ai.djl.nn.AbstractBaseBlock.forward(AbstractBaseBlock.java:93)
	at ai.djl.nn.SequentialBlock.forwardInternal(SequentialBlock.java:211)
	at ai.djl.nn.AbstractBaseBlock.forward(AbstractBaseBlock.java:93)
	at ai.djl.training.Trainer.forward(Trainer.java:187)
	at ai.djl.training.EasyTrain.trainSplit(EasyTrain.java:122)
	at ai.djl.training.EasyTrain.trainBatch(EasyTrain.java:110)
	at ai.djl.training.EasyTrain.fit(EasyTrain.java:58)
	at ai.djl.pytorch.integration.MpsTest.main(MpsTest.java:114)
```
because [`ParameterStore.deviceMap`](https://github.com/deepjavalibrary/djl/blob/50f8af6db5c55c3ae05faf8273523e5c2761a430/api/src/main/java/ai/djl/training/ParameterStore.java#L104) expects the value to be `0`.
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1365 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.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) 72.08% compared to head [(`a5c9623`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.22%.
> Report is 921 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...va/ai/djl/modality/nlp/generate/TextGenerator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9UZXh0R2VuZXJhdG9yLmphdmE=) | 2.81% | [276 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/nlp/generate/SeqBatcher.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaGVyLmphdmE=) | 0.75% | [132 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...ity/nlp/generate/ContrastiveSeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9Db250cmFzdGl2ZVNlcUJhdGNoU2NoZWR1bGVyLmphdmE=) | 2.97% | [98 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/djl/modality/nlp/generate/SeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaFNjaGVkdWxlci5qYXZh) | 9.83% | [55 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/cv/BufferedImageFactory.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0J1ZmZlcmVkSW1hZ2VGYWN0b3J5LmphdmE=) | 40.96% | [47 Missing and 2 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...a/ai/djl/modality/nlp/generate/StepGeneration.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TdGVwR2VuZXJhdGlvbi5qYXZh) | 2.04% | [48 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [api/src/main/java/ai/djl/ndarray/NDArray.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheS5qYXZh) | 43.42% | [39 Missing and 4 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...n/java/ai/djl/modality/cv/output/CategoryMask.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9DYXRlZ29yeU1hc2suamF2YQ==) | 22.00% | [39 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheUFkYXB0ZXIuamF2YQ==) | 71.21% | [31 Missing and 7 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | 37.50% | [35 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| ... and [76 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

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


```diff
@@             Coverage Diff              @@
##             master    #2873      +/-   ##
============================================
+ Coverage     72.08%   72.22%   +0.13%     
- Complexity     5126     7161    +2035     
============================================
  Files           473      708     +235     
  Lines         21970    31922    +9952     
  Branches       2351     3317     +966     
============================================
+ Hits          15838    23057    +7219     
- Misses         4925     7289    +2364     
- Partials       1207     1576     +369     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2873?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](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/.github/workflows/native_s3_huggingface.yml b/.github/workflows/native_s3_huggingface.yml
index e5a3d57fe..0492ce67e 100644
--- a/.github/workflows/native_s3_huggingface.yml
+++ b/.github/workflows/native_s3_huggingface.yml
@@ -155,8 +155,8 @@ jobs:
       - name: Copy files to S3 with the AWS CLI
         run: |
           TOKENIZERS_VERSION="$(cat gradle.properties | awk -F '=' '/tokenizers_version/ {print $2}')"
-          aws s3 sync extensions/tokenizers/jnilib s3://djl-ai/publish/tokenizers/${TOKENIZERS_VERSION}/jnilib/
-          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/tokenizers/${TOKENIZERS_VERSION}/jnilib/*"
+          /opt/homebrew/bin/aws s3 sync extensions/tokenizers/jnilib s3://djl-ai/publish/tokenizers/${TOKENIZERS_VERSION}/jnilib/
+          /opt/homebrew/bin/aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/tokenizers/${TOKENIZERS_VERSION}/jnilib/*"
 
   create-aarch64-runner:
     if: github.repository == 'deepjavalibrary/djl'
@@ -179,15 +179,14 @@ jobs:
     runs-on: [ self-hosted, aarch64 ]
     timeout-minutes: 30
     needs: create-aarch64-runner
-    container: centos:centos7
+    container: amazonlinux:2
     steps:
       - name: Install Environment
         run: |
           yum -y update
-          yum -y install centos-release-scl-rh epel-release perl-core
-          yum -y install devtoolset-7 git patch cmake3 libstdc++-static
+          yum -y groupinstall "Development Tools"
+          yum -y install patch perl-IPC-Cmd cmake3
           ln -s /usr/bin/cmake3 /usr/bin/cmake
-          curl https://sh.rustup.rs -sSf | sh -s -- -y
           pip3 install awscli --upgrade
       - uses: actions-rs/toolchain@v1
         with:
@@ -207,8 +206,6 @@ jobs:
             ${{ runner.os }}-gradle-
       - name: Release JNI prep
         run: |
-          source "$HOME/.cargo/env"
-          export PATH=$PATH:/opt/rh/devtoolset-7/root/usr/bin
           ./gradlew :extensions:tokenizers:compileJNI
           PYTORCH_PRECXX11=true ./gradlew -Pjni :extensions:tokenizers:test
       - name: Configure AWS Credentials
diff --git a/api/src/main/java/ai/djl/modality/Classifications.java b/api/src/main/java/ai/djl/modality/Classifications.java
index 84025ce07..13f3325b1 100644
--- a/api/src/main/java/ai/djl/modality/Classifications.java
+++ b/api/src/main/java/ai/djl/modality/Classifications.java
@@ -88,10 +88,18 @@ public class Classifications implements JsonSerializable, Ensembleable<Classific
      */
     public Classifications(List<String> classNames, NDArray probabilities, int topK) {
         this.classNames = classNames;
-        NDArray array = probabilities.toType(DataType.FLOAT64, false);
-        this.probabilities =
-                Arrays.stream(array.toDoubleArray()).boxed().collect(Collectors.toList());
-        array.close();
+        if (probabilities.getDataType() == DataType.FLOAT32) {
+            // Avoid converting float32 to float64 as this is not supported on MPS device
+            this.probabilities = new ArrayList<>();
+            for (float prob : probabilities.toFloatArray()) {
+                this.probabilities.add((double) prob);
+            }
+        } else {
+            NDArray array = probabilities.toType(DataType.FLOAT64, false);
+            this.probabilities =
+                    Arrays.stream(array.toDoubleArray()).boxed().collect(Collectors.toList());
+            array.close();
+        }
         this.topK = topK;
     }
 
diff --git a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
index b12ac5dd0..07e56a5ca 100644
--- a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
+++ b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
@@ -434,7 +434,12 @@ public interface NDArrayEx {
             if (dim == 3) {
                 result = result.expandDims(0);
             }
-            result = result.div(255.0).transpose(0, 3, 1, 2);
+            // For Apple Silicon MPS it is important not to switch to 64-bit float here
+            if (result.getDataType() == DataType.FLOAT32) {
+                result = result.div(255.0f).transpose(0, 3, 1, 2);
+            } else {
+                result = result.div(255.0).transpose(0, 3, 1, 2);
+            }
             if (dim == 3) {
                 result = result.squeeze(0);
             }
diff --git a/api/src/main/java/ai/djl/repository/AbstractRepository.java b/api/src/main/java/ai/djl/repository/AbstractRepository.java
index 3b83c359a..2ca087fdf 100644
--- a/api/src/main/java/ai/djl/repository/AbstractRepository.java
+++ b/api/src/main/java/ai/djl/repository/AbstractRepository.java
@@ -243,7 +243,7 @@ public abstract class AbstractRepository implements Repository {
         bis = new CloseShieldFilterInputStream(bis);
         try (TarArchiveInputStream tis = new TarArchiveInputStream(bis)) {
             TarArchiveEntry entry;
-            while ((entry = tis.getNextTarEntry()) != null) {
+            while ((entry = tis.getNextEntry()) != null) {
                 String entryName = entry.getName();
                 if (entryName.contains("..")) {
                     throw new IOException("Malicious zip entry: " + entryName);
diff --git a/api/src/main/java/ai/djl/repository/RemoteRepository.java b/api/src/main/java/ai/djl/repository/RemoteRepository.java
index 6b01ce14e..52f87afbf 100644
--- a/api/src/main/java/ai/djl/repository/RemoteRepository.java
+++ b/api/src/main/java/ai/djl/repository/RemoteRepository.java
@@ -75,7 +75,7 @@ public class RemoteRepository extends AbstractRepository {
                 Metadata metadata = JsonUtils.GSON_PRETTY.fromJson(reader, Metadata.class);
                 metadata.init(arguments);
                 Date lastUpdated = metadata.getLastUpdated();
-                if (Boolean.getBoolean("offline")
+                if (Utils.isOfflineMode()
                         || System.currentTimeMillis() - lastUpdated.getTime() < ONE_DAY) {
                     metadata.setRepositoryUri(mrlUri);
                     return metadata;
diff --git a/api/src/main/java/ai/djl/util/Ec2Utils.java b/api/src/main/java/ai/djl/util/Ec2Utils.java
index 178c3d7ef..540818296 100644
--- a/api/src/main/java/ai/djl/util/Ec2Utils.java
+++ b/api/src/main/java/ai/djl/util/Ec2Utils.java
@@ -97,7 +97,7 @@ public final class Ec2Utils {
      * @param engine the default engine name
      */
     public static void callHome(String engine) {
-        if (Boolean.getBoolean("offline")
+        if (Utils.isOfflineMode()
                 || Boolean.parseBoolean(Utils.getEnvOrSystemProperty("OPT_OUT_TRACKING"))
                 || System.currentTimeMillis() - lastCheckIn < ONE_DAY) {
             return;
diff --git a/api/src/main/java/ai/djl/util/Utils.java b/api/src/main/java/ai/djl/util/Utils.java
index c8e1bd514..270958d5b 100644
--- a/api/src/main/java/ai/djl/util/Utils.java
+++ b/api/src/main/java/ai/djl/util/Utils.java
@@ -357,6 +357,20 @@ public final class Utils {
         return Paths.get(cacheDir);
     }
 
+    /**
+     * Returns if offline mode is enabled.
+     *
+     * @return true if offline mode is enabled
+     */
+    public static boolean isOfflineMode() {
+        String mode = getenv("DJL_OFFLINE", System.getProperty("ai.djl.offline"));
+        if (mode != null) {
+            return Boolean.parseBoolean(mode);
+        }
+        // backward compatible
+        return Boolean.getBoolean("offline");
+    }
+
     /**
      * Returns nested model directory if the directory contains only one subdirectory.
      *
@@ -481,7 +495,7 @@ public final class Utils {
      */
     public static InputStream openUrl(URL url) throws IOException {
         String protocol = url.getProtocol();
-        if (Boolean.getBoolean("offline")
+        if (isOfflineMode()
                 && ("http".equalsIgnoreCase(protocol) || "https".equalsIgnoreCase(protocol))) {
             throw new IOException("Offline model is enabled.");
         }
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java
index a92a9b6a3..deef04907 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java
@@ -30,6 +30,7 @@ import ai.djl.util.Utils;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
 import java.util.Map;
 
 /**
@@ -118,8 +119,9 @@ public final class FashionMnist extends ArrayDataset {
             byte[] buf = Utils.toByteArray(is);
             try (NDArray array =
                     manager.create(
-                            new Shape(length, IMAGE_WIDTH, IMAGE_HEIGHT, 1), DataType.UINT8)) {
-                array.set(buf);
+                            ByteBuffer.wrap(buf),
+                            new Shape(length, IMAGE_WIDTH, IMAGE_HEIGHT, 1),
+                            DataType.UINT8)) {
                 return array.toType(DataType.FLOAT32, false);
             }
         }
@@ -132,8 +134,8 @@ public final class FashionMnist extends ArrayDataset {
             }
 
             byte[] buf = Utils.toByteArray(is);
-            try (NDArray array = manager.create(new Shape(buf.length), DataType.UINT8)) {
-                array.set(buf);
+            try (NDArray array =
+                    manager.create(ByteBuffer.wrap(buf), new Shape(buf.length), DataType.UINT8)) {
                 return array.toType(DataType.FLOAT32, false);
             }
         }
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java
index 164ba9876..5503e721c 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java
@@ -30,6 +30,7 @@ import ai.djl.util.Utils;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
 import java.util.Map;
 
 /**
@@ -111,8 +112,9 @@ public final class Mnist extends ArrayDataset {
             }
 
             byte[] buf = Utils.toByteArray(is);
-            try (NDArray array = manager.create(new Shape(length, 28, 28, 1), DataType.UINT8)) {
-                array.set(buf);
+            try (NDArray array =
+                    manager.create(
+                            ByteBuffer.wrap(buf), new Shape(length, 28, 28, 1), DataType.UINT8)) {
                 return array.toType(DataType.FLOAT32, false);
             }
         }
@@ -123,10 +125,9 @@ public final class Mnist extends ArrayDataset {
             if (is.skip(8) != 8) {
                 throw new AssertionError("Failed skip data.");
             }
-
             byte[] buf = Utils.toByteArray(is);
-            try (NDArray array = manager.create(new Shape(buf.length), DataType.UINT8)) {
-                array.set(buf);
+            try (NDArray array =
+                    manager.create(ByteBuffer.wrap(buf), new Shape(buf.length), DataType.UINT8)) {
                 return array.toType(DataType.FLOAT32, false);
             }
         }
diff --git a/bom/build.gradle b/bom/build.gradle
index 736e8e4a5..aa17702f9 100644
--- a/bom/build.gradle
+++ b/bom/build.gradle
@@ -109,15 +109,15 @@ publishing {
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-mkl", "linux-x86_64", "${mxnet_version}")
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-mkl", "win-x86_64", "${mxnet_version}")
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-cu112mkl", "linux-x86_64", "${mxnet_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-x86_64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-aarch64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "win-x86_64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-aarch64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "win-x86_64", "${pytorch_version}-SNAPSHOT")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121-precxx11", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-aarch64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "linux-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "win-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-aarch64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "linux-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "win-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121-precxx11", "linux-x86_64", "${pytorch_version}")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117", "linux-x86_64", "1.13.1")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117", "win-x86_64", "1.13.1")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117-precxx11", "linux-x86_64", "1.13.1")
diff --git a/build.gradle b/build.gradle
index f98b86c4e..4b8b20cd5 100644
--- a/build.gradle
+++ b/build.gradle
@@ -88,7 +88,7 @@ configure(javaProjects()) {
             systemProperty "disableProgressBar", "true"
             systemProperty "nightly", System.getProperty("nightly", "false")
             if (gradle.startParameter.offline) {
-                systemProperty "offline", "true"
+                systemProperty "ai.djl.offline", "true"
             }
             // This is used to avoid overriding on default engine for modules:
             // mxnet-engine, mxnet-model-zoo, api (MockEngine), basicdataset, fasttext, etc
diff --git a/docs/telemetry.md b/docs/telemetry.md
index d6ff9b20b..256adf00a 100644
--- a/docs/telemetry.md
+++ b/docs/telemetry.md
@@ -20,5 +20,5 @@ System.setProperty("OPT_OUT_TRACKING", "true")
 Usage tracking is also disable in `offline` mode:
 
 ```java
-System.setProperty("offline", "true")
+System.setProperty("ai.djl.offline", "true")
 ```
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngineProvider.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngineProvider.java
index 1b9cdd0ab..42ca3c5b8 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngineProvider.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngineProvider.java
@@ -18,6 +18,8 @@ import ai.djl.engine.EngineProvider;
 /** {@code PtEngineProvider} is the PyTorch implementation of {@link EngineProvider}. */
 public class PtEngineProvider implements EngineProvider {
 
+    private static volatile Engine engine; // NOPMD
+
     /** {@inheritDoc} */
     @Override
     public String getEngineName() {
@@ -33,10 +35,13 @@ public class PtEngineProvider implements EngineProvider {
     /** {@inheritDoc} */
     @Override
     public Engine getEngine() {
-        return InstanceHolder.INSTANCE;
-    }
-
-    private static class InstanceHolder {
-        static final Engine INSTANCE = PtEngine.newInstance();
+        if (engine == null) {
+            synchronized (PtEngineProvider.class) {
+                if (engine == null) {
+                    engine = PtEngine.newInstance();
+                }
+            }
+        }
+        return engine;
     }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
index 9d4224639..860c1bc48 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
@@ -106,10 +106,6 @@ public final class LibUtils {
 
     private static void loadLibTorch(LibTorch libTorch) {
         Path libDir = libTorch.dir.toAbsolutePath();
-        if ("1.8.1".equals(getVersion()) && System.getProperty("os.name").startsWith("Mac")) {
-            // PyTorch 1.8.1 libtorch_cpu.dylib cannot be loaded individually
-            return;
-        }
         boolean isCuda = libTorch.flavor.contains("cu");
         List<String> deferred =
                 Arrays.asList(
@@ -120,6 +116,7 @@ public final class LibUtils {
                         System.mapLibraryName("torch_cuda_cpp"),
                         System.mapLibraryName("torch_cuda_cu"),
                         System.mapLibraryName("torch_cuda"),
+                        System.mapLibraryName("nvfuser_codegen"),
                         System.mapLibraryName("torch"));
 
         Set<String> loadLater = new HashSet<>(deferred);
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java
index 9ee8fc19c..a0099073c 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java
@@ -131,8 +131,7 @@ public class HfModelZoo extends ModelZoo {
             Path file = dir.resolve("models.json");
             if (Files.exists(file)) {
                 long lastModified = Files.getLastModifiedTime(file).toMillis();
-                if (Boolean.getBoolean("offline")
-                        || System.currentTimeMillis() - lastModified < ONE_DAY) {
+                if (Utils.isOfflineMode() || System.currentTimeMillis() - lastModified < ONE_DAY) {
                     try (Reader reader = Files.newBufferedReader(file)) {
                         return JsonUtils.GSON.fromJson(reader, type);
                     }
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index f3b85c241..1efabb61c 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -41,9 +41,20 @@ class HuggingfaceConverter:
         if not os.path.exists(temp_dir):
             os.makedirs(temp_dir)
 
-        hf_pipeline = self.load_model(model_id)
-        # Save tokenizer.json to temp dir
-        self.save_tokenizer(hf_pipeline, temp_dir)
+        try:
+            hf_pipeline = self.load_model(model_id)
+        except Exception as e:
+            logging.warning(f"Failed to load model: {model_id}.")
+            logging.warning(e, exc_info=True)
+            return False, "Failed to load model", -1
+
+        try:
+            # Save tokenizer.json to temp dir
+            self.save_tokenizer(hf_pipeline, temp_dir)
+        except Exception as e:
+            logging.warning(f"Failed to save tokenizer: {model_id}.")
+            logging.warning(e, exc_info=True)
+            return False, "Failed to save tokenizer", -1
 
         # Save config.json just for reference
         config = hf_hub_download(repo_id=model_id, filename="config.json")
@@ -112,7 +123,7 @@ class HuggingfaceConverter:
             logging.info(f"Saving torchscript model: {model_name}.pt ...")
             model_file = os.path.join(temp_dir, f"{model_name}.pt")
             script_module.save(model_file)
-        except (RuntimeError, ValueError) as e:
+        except Exception as e:
             logging.warning(f"Failed to trace model: {model_id}.")
             logging.warning(e, exc_info=True)
             return None
diff --git a/extensions/tokenizers/src/main/python/huggingface_models.py b/extensions/tokenizers/src/main/python/huggingface_models.py
index 549db3788..12db9cf2c 100644
--- a/extensions/tokenizers/src/main/python/huggingface_models.py
+++ b/extensions/tokenizers/src/main/python/huggingface_models.py
@@ -56,6 +56,8 @@ class HuggingfaceModels:
         self.temp_dir = f"{self.output_dir}/tmp"
 
     def list_models(self, args: Namespace) -> List[dict]:
+        import_all = os.environ.get("HF_IMPORT_ALL")
+
         api = HfApi()
         if args.model_name:
             models = api.list_models(filter="pytorch",
@@ -63,16 +65,20 @@ class HuggingfaceModels:
                                      sort="downloads",
                                      direction=-1,
                                      limit=args.limit)
-            if not models:
-                logging.warning(f"no model found: {args.model_name}.")
+            import_all = True
         else:
             models = api.list_models(filter=f"{args.category},pytorch",
                                      sort="downloads",
                                      direction=-1,
                                      limit=args.limit)
-            if not models:
+        if not models:
+            if args.model_name:
+                logging.warning(f"no model found: {args.model_name}.")
+            else:
                 logging.warning(f"no model matches category: {args.category}.")
 
+            return []
+
         ret = []
         for model_info in models:
             model_id = model_info.modelId
@@ -83,7 +89,7 @@ class HuggingfaceModels:
                 continue
 
             languages = get_lang_tags(model_info)
-            if "en" not in languages and not os.environ["HF_IMPORT_ALL"]:
+            if "en" not in languages and not import_all:
                 logging.warning(f"Skip non-English model: {model_id}.")
                 continue
 
@@ -94,6 +100,12 @@ class HuggingfaceModels:
                     logging.info(f"Skip converted model: {model_id}.")
                     continue
 
+            if model_info.downloads < 50 and not import_all:
+                logging.info(
+                    f"Skip model {model_info.modelId}, downloads {model_info.downloads} < 50"
+                )
+                continue
+
             try:
                 config = hf_hub_download(repo_id=model_id,
                                          filename="config.json")
diff --git a/extensions/tokenizers/src/main/python/model_zoo_importer.py b/extensions/tokenizers/src/main/python/model_zoo_importer.py
index 9ed32ec58..0ed67bd10 100644
--- a/extensions/tokenizers/src/main/python/model_zoo_importer.py
+++ b/extensions/tokenizers/src/main/python/model_zoo_importer.py
@@ -49,9 +49,17 @@ def main():
         model_info = model["model_info"]
         converter = SUPPORTED_TASK[task]
 
-        result, reason, size = converter.save_model(model_info, args, temp_dir)
-        if not result:
-            logging.error(f"{model_info.modelId}: {reason}")
+        try:
+            result, reason, size = converter.save_model(
+                model_info, args, temp_dir)
+            if not result:
+                logging.error(f"{model_info.modelId}: {reason}")
+        except Exception as e:
+            logging.warning(f"Failed to convert model: {model_info.modelId}.")
+            logging.warning(e, exc_info=True)
+            result = False
+            reason = "Failed to convert model"
+            size = -1
 
         huggingface_models.update_progress(model_info, converter.application,
                                            result, reason, size, args.cpu_only)
diff --git a/gradle.properties b/gradle.properties
index 20426a4d0..9c8822596 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -20,20 +20,20 @@ trt_version=8.4.1
 onnxruntime_version=1.16.0
 paddlepaddle_version=2.3.2
 sentencepiece_version=0.1.97
-tokenizers_version=0.14.1
+tokenizers_version=0.15.0
 fasttext_version=0.9.2
 xgboost_version=2.0.1
 lightgbm_version=3.2.110
 rapis_version=22.12.0
 
-commons_cli_version=1.5.0
-commons_compress_version=1.23.0
+commons_cli_version=1.6.0
+commons_compress_version=1.25.0
 commons_csv_version=1.10.0
 commons_logging_version=1.2
 gson_version=2.10.1
 jna_version=5.13.0
 slf4j_version=1.7.36
-log4j_slf4j_version=2.21.0
+log4j_slf4j_version=2.22.0
 awssdk_version=2.20.121
 hadoop_version=3.3.5
 javacpp_version=1.5.9


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
index 8b4e2326f..5b6ed349e 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
@@ -13,6 +13,7 @@
 package ai.djl.pytorch.integration;
 
 import ai.djl.Device;
+import ai.djl.modality.Classifications;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
@@ -21,6 +22,10 @@ import org.testng.Assert;
 import org.testng.SkipException;
 import org.testng.annotations.Test;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
 public class MpsTest {
 
     @Test
@@ -36,4 +41,39 @@ public class MpsTest {
             Assert.assertEquals(array.getDevice().getDeviceType(), "mps");
         }
     }
+
+    private static boolean checkMpsCompatible() {
+        return "aarch64".equals(System.getProperty("os.arch"))
+                && System.getProperty("os.name").startsWith("Mac");
+    }
+
+    @Test
+    public void testToTensorMPS() {
+        if (!checkMpsCompatible()) {
+            throw new SkipException("MPS toTensor test requires Apple Silicon macOS.");
+        }
+
+        // Test that toTensor does not fail on MPS (e.g. due to use of float64 for division)
+        try (NDManager manager = NDManager.newBaseManager(Device.fromName("mps"))) {
+            NDArray array = manager.create(127f).reshape(1, 1, 1, 1);
+            NDArray tensor = array.getNDArrayInternal().toTensor();
+            Assert.assertEquals(tensor.toFloatArray(), new float[] {127f / 255f});
+        }
+    }
+
+    @Test
+    public void testClassificationsMPS() {
+        if (!checkMpsCompatible()) {
+            throw new SkipException("MPS classification test requires Apple Silicon macOS.");
+        }
+
+        // Test that classifications do not fail on MPS (e.g. due to conversion of probabilities to
+        // float64)
+        try (NDManager manager = NDManager.newBaseManager(Device.fromName("mps"))) {
+            List<String> names = Arrays.asList("First", "Second", "Third", "Fourth", "Fifth");
+            NDArray tensor = manager.create(new float[] {0f, 0.125f, 1f, 0.5f, 0.25f});
+            Classifications classifications = new Classifications(names, tensor);
+            Assert.assertEquals(classifications.topK(1), Collections.singletonList("Third"));
+        }
+    }
 }
diff --git a/examples/src/test/java/ai/djl/testing/TestRequirements.java b/examples/src/test/java/ai/djl/testing/TestRequirements.java
index e8c9bd4bd..01eef7562 100644
--- a/examples/src/test/java/ai/djl/testing/TestRequirements.java
+++ b/examples/src/test/java/ai/djl/testing/TestRequirements.java
@@ -14,6 +14,7 @@ package ai.djl.testing;
 
 import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
+import ai.djl.util.Utils;
 
 import org.testng.SkipException;
 
@@ -45,7 +46,7 @@ public final class TestRequirements {
 
     /** Requires a test not be run in offline mode. */
     public static void notOffline() {
-        if (Boolean.getBoolean("offline")) {
+        if (Utils.isOfflineMode()) {
             throw new SkipException("This test can not run while offline");
         }
     }
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java
index e585a219f..5c3ed6c3e 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java
@@ -103,21 +103,21 @@ public class ModelZooTest {
     @Test
     public void testOffLine() throws IOException {
         System.setProperty("DJL_CACHE_DIR", "build/cache");
-        System.setProperty("offline", "true");
+        System.setProperty("ai.djl.offline", "true");
         try {
             Utils.deleteQuietly(Paths.get("build/cache"));
             // static variables cannot not be initialized properly if directly use new HfModelZoo()
             ModelZoo.getModelZoo("ai.djl.huggingface.pytorch");
 
             ModelZoo zoo = new HfModelZoo();
-            Assert.assertTrue(zoo.getModelLoaders().size() > 0);
+            Assert.assertFalse(zoo.getModelLoaders().isEmpty());
 
             Set<String> engines = zoo.getSupportedEngines();
             Assert.assertEquals(engines.size(), 1);
             Assert.assertEquals(engines.iterator().next(), "PyTorch");
         } finally {
             System.clearProperty("DJL_CACHE_DIR");
-            System.clearProperty("offline");
+            System.clearProperty("ai.djl.offline");
         }
     }
 }
diff --git a/testing/src/main/java/ai/djl/testing/TestRequirements.java b/testing/src/main/java/ai/djl/testing/TestRequirements.java
index bf57d64bd..32f242589 100644
--- a/testing/src/main/java/ai/djl/testing/TestRequirements.java
+++ b/testing/src/main/java/ai/djl/testing/TestRequirements.java
@@ -13,6 +13,7 @@
 package ai.djl.testing;
 
 import ai.djl.engine.Engine;
+import ai.djl.util.Utils;
 
 import org.testng.SkipException;
 
@@ -45,7 +46,7 @@ public final class TestRequirements {
 
     /** Requires a test not be run in offline mode. */
     public static void notOffline() {
-        if (Boolean.getBoolean("offline")) {
+        if (Utils.isOfflineMode()) {
             throw new SkipException("This test can not run while offline");
         }
     }