From e12d200c97d7aab668b976e92b46513c9ca7a0d8 Mon Sep 17 00:00:00 2001
From: Aleksei Nikiforov
 <103434461+AlekseiNikiforovIBM@users.noreply.github.com>
Date: Wed, 6 Dec 2023 17:18:56 +0100
Subject: [PATCH] S390x big endian fixes (#8149)

Fixes for multiple tests on s390x
---
 test/test_transforms.py              |  9 +++++----
 torchvision/datasets/mnist.py        | 16 +++++++++++++---
 torchvision/transforms/functional.py |  5 +++--
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/test/test_transforms.py b/test/test_transforms.py
index 4440942f7cf..01c27a30a65 100644
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -2,6 +2,7 @@
 import os
 import random
 import re
+import sys
 from functools import partial
 
 import numpy as np
@@ -614,7 +615,7 @@ def _get_1_channel_tensor_various_types():
 
         img_data_short = torch.ShortTensor(1, 4, 4).random_()
         expected_output = img_data_short.numpy()
-        yield img_data_short, expected_output, "I;16"
+        yield img_data_short, expected_output, "I;16" if sys.byteorder == "little" else "I;16B"
 
         img_data_int = torch.IntTensor(1, 4, 4).random_()
         expected_output = img_data_int.numpy()
@@ -631,7 +632,7 @@ def _get_2d_tensor_various_types():
 
         img_data_short = torch.ShortTensor(4, 4).random_()
         expected_output = img_data_short.numpy()
-        yield img_data_short, expected_output, "I;16"
+        yield img_data_short, expected_output, "I;16" if sys.byteorder == "little" else "I;16B"
 
         img_data_int = torch.IntTensor(4, 4).random_()
         expected_output = img_data_int.numpy()
@@ -662,7 +663,7 @@ def test_1_channel_float_tensor_to_pil_image(self):
         [
             (torch.Tensor(4, 4, 1).uniform_().numpy(), "L"),
             (torch.ByteTensor(4, 4, 1).random_(0, 255).numpy(), "L"),
-            (torch.ShortTensor(4, 4, 1).random_().numpy(), "I;16"),
+            (torch.ShortTensor(4, 4, 1).random_().numpy(), "I;16" if sys.byteorder == "little" else "I;16B"),
             (torch.IntTensor(4, 4, 1).random_().numpy(), "I"),
         ],
     )
@@ -744,7 +745,7 @@ def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expe
         [
             (torch.Tensor(4, 4).uniform_().numpy(), "L"),
             (torch.ByteTensor(4, 4).random_(0, 255).numpy(), "L"),
-            (torch.ShortTensor(4, 4).random_().numpy(), "I;16"),
+            (torch.ShortTensor(4, 4).random_().numpy(), "I;16" if sys.byteorder == "little" else "I;16B"),
             (torch.IntTensor(4, 4).random_().numpy(), "I"),
         ],
     )
diff --git a/torchvision/datasets/mnist.py b/torchvision/datasets/mnist.py
index 6953d1fc5c2..e5d74ad6c1d 100644
--- a/torchvision/datasets/mnist.py
+++ b/torchvision/datasets/mnist.py
@@ -510,15 +510,25 @@ def read_sn3_pascalvincent_tensor(path: str, strict: bool = True) -> torch.Tenso
     # read
     with open(path, "rb") as f:
         data = f.read()
+
     # parse
-    magic = get_int(data[0:4])
-    nd = magic % 256
-    ty = magic // 256
+    if sys.byteorder == "little":
+        magic = get_int(data[0:4])
+        nd = magic % 256
+        ty = magic // 256
+    else:
+        nd = get_int(data[0:1])
+        ty = get_int(data[1:2]) + get_int(data[2:3]) * 256 + get_int(data[3:4]) * 256 * 256
+
     assert 1 <= nd <= 3
     assert 8 <= ty <= 14
     torch_type = SN3_PASCALVINCENT_TYPEMAP[ty]
     s = [get_int(data[4 * (i + 1) : 4 * (i + 2)]) for i in range(nd)]
 
+    if sys.byteorder == "big":
+        for i in range(len(s)):
+            s[i] = int.from_bytes(s[i].to_bytes(4, byteorder="little"), byteorder="big", signed=False)
+
     parsed = torch.frombuffer(bytearray(data), dtype=torch_type, offset=(4 * (nd + 1)))
 
     # The MNIST format uses the big endian byte order, while `torch.frombuffer` uses whatever the system uses. In case
diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py
index 7cbe2d99071..698942c56af 100644
--- a/torchvision/transforms/functional.py
+++ b/torchvision/transforms/functional.py
@@ -1,5 +1,6 @@
 import math
 import numbers
+import sys
 import warnings
 from enum import Enum
 from typing import Any, List, Optional, Tuple, Union
@@ -162,7 +163,7 @@ def to_tensor(pic) -> Tensor:
         return torch.from_numpy(nppic).to(dtype=default_float_dtype)
 
     # handle PIL Image
-    mode_to_nptype = {"I": np.int32, "I;16": np.int16, "F": np.float32}
+    mode_to_nptype = {"I": np.int32, "I;16" if sys.byteorder == "little" else "I;16B": np.int16, "F": np.float32}
     img = torch.from_numpy(np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True))
 
     if pic.mode == "1":
@@ -285,7 +286,7 @@ def to_pil_image(pic, mode=None):
         if npimg.dtype == np.uint8:
             expected_mode = "L"
         elif npimg.dtype == np.int16:
-            expected_mode = "I;16"
+            expected_mode = "I;16" if sys.byteorder == "little" else "I;16B"
         elif npimg.dtype == np.int32:
             expected_mode = "I"
         elif npimg.dtype == np.float32: