Refactor conversion function (#19799)

* Refactor conversion function * Remove dupe line * Fixes * Fixes * Use the right variable... * Fix last test
2022-10-24 13:48:40 -04:00
parent 9ecb13d63a
commit d4eb52d13d
4 changed files with 316 additions and 75 deletions
--- a/tests/utils/test_generic.py
+++ b/tests/utils/test_generic.py
@@ -15,7 +15,29 @@

 import unittest

-from transformers.utils import flatten_dict
+import numpy as np
+
+from transformers.testing_utils import require_flax, require_tf, require_torch
+from transformers.utils import (
+    expand_dims,
+    flatten_dict,
+    is_flax_available,
+    is_tf_available,
+    is_torch_available,
+    reshape,
+    squeeze,
+    transpose,
+)
+
+
+if is_flax_available():
+    import jax.numpy as jnp
+
+if is_tf_available():
+    import tensorflow as tf
+
+if is_torch_available():
+    import torch


 class GenericTester(unittest.TestCase):
@@ -43,3 +65,136 @@ class GenericTester(unittest.TestCase):
        }

        self.assertEqual(flatten_dict(input_dict), expected_dict)
+
+    def test_transpose_numpy(self):
+        x = np.random.randn(3, 4)
+        self.assertTrue(np.allclose(transpose(x), x.transpose()))
+
+        x = np.random.randn(3, 4, 5)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), x.transpose((1, 2, 0))))
+
+    @require_torch
+    def test_transpose_torch(self):
+        x = np.random.randn(3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(transpose(x), transpose(t).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), transpose(t, axes=(1, 2, 0)).numpy()))
+
+    @require_tf
+    def test_transpose_tf(self):
+        x = np.random.randn(3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(transpose(x), transpose(t).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), transpose(t, axes=(1, 2, 0)).numpy()))
+
+    @require_flax
+    def test_transpose_flax(self):
+        x = np.random.randn(3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(transpose(x), np.asarray(transpose(t))))
+
+        x = np.random.randn(3, 4, 5)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), np.asarray(transpose(t, axes=(1, 2, 0)))))
+
+    def test_reshape_numpy(self):
+        x = np.random.randn(3, 4)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), np.reshape(x, (4, 3))))
+
+        x = np.random.randn(3, 4, 5)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), np.reshape(x, (12, 5))))
+
+    @require_torch
+    def test_reshape_torch(self):
+        x = np.random.randn(3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), reshape(t, (4, 3)).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), reshape(t, (12, 5)).numpy()))
+
+    @require_tf
+    def test_reshape_tf(self):
+        x = np.random.randn(3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), reshape(t, (4, 3)).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), reshape(t, (12, 5)).numpy()))
+
+    @require_flax
+    def test_reshape_flax(self):
+        x = np.random.randn(3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), np.asarray(reshape(t, (4, 3)))))
+
+        x = np.random.randn(3, 4, 5)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), np.asarray(reshape(t, (12, 5)))))
+
+    def test_squeeze_numpy(self):
+        x = np.random.randn(1, 3, 4)
+        self.assertTrue(np.allclose(squeeze(x), np.squeeze(x)))
+
+        x = np.random.randn(1, 4, 1, 5)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), np.squeeze(x, axis=2)))
+
+    @require_torch
+    def test_squeeze_torch(self):
+        x = np.random.randn(1, 3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(squeeze(x), squeeze(t).numpy()))
+
+        x = np.random.randn(1, 4, 1, 5)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), squeeze(t, axis=2).numpy()))
+
+    @require_tf
+    def test_squeeze_tf(self):
+        x = np.random.randn(1, 3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(squeeze(x), squeeze(t).numpy()))
+
+        x = np.random.randn(1, 4, 1, 5)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), squeeze(t, axis=2).numpy()))
+
+    @require_flax
+    def test_squeeze_flax(self):
+        x = np.random.randn(1, 3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(squeeze(x), np.asarray(squeeze(t))))
+
+        x = np.random.randn(1, 4, 1, 5)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), np.asarray(squeeze(t, axis=2))))
+
+    def test_expand_dims_numpy(self):
+        x = np.random.randn(3, 4)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), np.expand_dims(x, axis=1)))
+
+    @require_torch
+    def test_expand_dims_torch(self):
+        x = np.random.randn(3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), expand_dims(t, axis=1).numpy()))
+
+    @require_tf
+    def test_expand_dims_tf(self):
+        x = np.random.randn(3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), expand_dims(t, axis=1).numpy()))
+
+    @require_flax
+    def test_expand_dims_flax(self):
+        x = np.random.randn(3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), np.asarray(expand_dims(t, axis=1))))