test(integration): add tests for integration

optree/integration/jax.py

@@ -159,7 +159,13 @@ def _unravel_pytree(
     return tree_unflatten(treespec, unravel_func(flat))
 
 
-def _unravel_empty(_: Array) -> list[ArrayLike]:
+def _unravel_empty(flat: Array) -> list[ArrayLike]:
+    if jnp.shape(flat) != (0,):
+        raise ValueError(
+            f'The unravel function expected an array of shape {(0,)}, '
+            f'got shape {jnp.shape(flat)}.',
+        )
+
     return []
 
 

optree/integration/numpy.py

@@ -88,7 +88,12 @@ def _unravel_pytree(
     return tree_unflatten(treespec, unravel_func(flat))
 
 
-def _unravel_empty(_: np.ndarray) -> list[np.ndarray]:
+def _unravel_empty(flat: np.ndarray) -> list[np.ndarray]:
+    if np.shape(flat) != (0,):
+        raise ValueError(
+            f'The unravel function expected an array of shape {(0,)}, '
+            f'got shape {np.shape(flat)}.',
+        )
     return []
 
 

optree/integration/torch.py

@@ -85,7 +85,13 @@ def _unravel_pytree(
     return tree_unflatten(treespec, unravel_func(flat))
 
 
-def _unravel_empty(_: torch.Tensor) -> list[torch.Tensor]:
+def _unravel_empty(flat: torch.Tensor) -> list[torch.Tensor]:
+    if not torch.is_tensor(flat):
+        raise ValueError(f'Expected a tensor to unravel, got {type(flat)}.')
+    if flat.shape != (0,):
+        raise ValueError(
+            f'The unravel function expected a tensor of shape {(0,)}, got shape {flat.shape}.',
+        )
     return []
 
 
@@ -126,9 +132,11 @@ def _unravel_leaves_single_dtype(
     shapes: tuple[tuple[int, ...]],
     flat: torch.Tensor,
 ) -> list[torch.Tensor]:
+    if not torch.is_tensor(flat):
+        raise ValueError(f'Expected a tensor to unravel, got {type(flat)}.')
     if flat.shape != (sum(sizes),):
         raise ValueError(
-            f'The unravel function expected an array of shape {(sum(sizes),)}, '
+            f'The unravel function expected a tensor of shape {(sum(sizes),)}, '
             f'got shape {flat.shape}.',
         )
 
@@ -147,7 +155,7 @@ def _unravel_leaves(
         raise ValueError(f'Expected a tensor to unravel, got {type(flat)}.')
     if flat.shape != (sum(sizes),):
         raise ValueError(
-            f'The unravel function expected an array of shape {(sum(sizes),)}, '
+            f'The unravel function expected a tensor of shape {(sum(sizes),)}, '
             f'got shape {flat.shape}.',
         )
     if flat.dtype != to_dtype:

tests/integration/test_jax.py

@@ -78,17 +78,16 @@ def replace_leaf(_):
         assert leaf.dtype == reconstructed_leaf.dtype
         assert leaf.shape == reconstructed_leaf.shape
 
-    if len(leaves) > 0:
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(flat.reshape((-1, 1)))
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(jnp.concatenate([flat, jnp.zeros((1,))]))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(flat.reshape((-1, 1)))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(jnp.concatenate([flat, jnp.zeros((1,))]))
 
     if all(dtypes.dtype(leaf) == dtypes.dtype(flat) for leaf in leaves):
         unravel_func(lax.convert_element_type(flat, jnp.complex128))
@@ -139,16 +138,15 @@ def replace_leaf(_):
         assert leaf.dtype == reconstructed_leaf.dtype
         assert leaf.shape == reconstructed_leaf.shape
 
-    if len(leaves) > 0:
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(flat.reshape((-1, 1)))
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(jnp.concatenate([flat, jnp.zeros((1,))]))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(flat.reshape((-1, 1)))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(jnp.concatenate([flat, jnp.zeros((1,))]))
 
     unravel_func(flat.astype(jnp.complex128))

tests/integration/test_numpy.py

@@ -72,17 +72,16 @@ def replace_leaf(_):
         assert leaf.dtype == reconstructed_leaf.dtype
         assert leaf.shape == reconstructed_leaf.shape
 
-    if len(leaves) > 0:
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(flat.reshape((-1, 1)))
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(np.concatenate([flat, np.zeros((1,))]))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(flat.reshape((-1, 1)))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(np.concatenate([flat, np.zeros((1,))]))
 
     if all(np.result_type(leaf) == flat.dtype for leaf in leaves):
         unravel_func(flat.astype(np.complex128))
@@ -132,16 +131,15 @@ def replace_leaf(_):
         assert leaf.dtype == reconstructed_leaf.dtype
         assert leaf.shape == reconstructed_leaf.shape
 
-    if len(leaves) > 0:
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(flat.reshape((-1, 1)))
-        with pytest.raises(
-            ValueError,
-            match=r'The unravel function expected an array of shape .*, got .*\.',
-        ):
-            unravel_func(np.concatenate([flat, np.zeros((1,))]))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(flat.reshape((-1, 1)))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected an array of shape .*, got .*\.',
+    ):
+        unravel_func(np.concatenate([flat, np.zeros((1,))]))
 
     unravel_func(flat.astype(np.complex128))

tests/integration/test_torch.py

@@ -0,0 +1,167 @@
+# Copyright 2022-2023 MetaOPT Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# pylint: disable=missing-function-docstring,wrong-import-position,wrong-import-order
+
+import pytest
+
+
+torch = pytest.importorskip('torch')
+
+import random
+
+import torch
+
+import optree
+from helpers import LEAVES, TREES, parametrize
+
+
+@parametrize(tree=list(TREES + LEAVES))
+def test_tree_ravel(tree):
+    random.seed(0)
+
+    def replace_leaf(_):
+        candidates = [
+            torch.tensor(random.randint(-100, 100)),
+            torch.tensor(random.uniform(-100.0, 100.0)),
+        ]
+
+        shapes = [
+            (),
+            (random.randint(1, 10),),
+            (random.randint(1, 10), random.randint(1, 10)),
+            (random.randint(1, 10), random.randint(1, 10), random.randint(1, 10)),
+        ]
+        dtypes = [
+            torch.float32,
+            torch.float64,
+            torch.int32,
+            torch.int64,
+        ]
+        for dtype in dtypes:
+            candidates.extend(
+                (5.0 * (2.0 * torch.randn(size=shape) - 1.0)).to(dtype) for shape in shapes
+            )
+
+        return random.choice(candidates)
+
+    tree = optree.tree_map(replace_leaf, tree)
+    flat, unravel_func = optree.integration.torch.tree_ravel(tree)
+
+    leaves, treespec = optree.tree_flatten(tree)
+    assert flat.numel() == sum(leaf.numel() for leaf in leaves)
+    assert flat.shape == (flat.numel(),)
+
+    reconstructed = unravel_func(flat)
+    reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed)
+    assert reconstructed_treespec == treespec
+    assert len(leaves) == len(reconstructed_leaves)
+    for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves):
+        assert torch.is_tensor(leaf)
+        assert torch.is_tensor(reconstructed_leaf)
+        assert torch.allclose(leaf, reconstructed_leaf)
+        assert leaf.dtype == reconstructed_leaf.dtype
+        assert leaf.shape == reconstructed_leaf.shape
+
+    with pytest.raises(ValueError, match=r'Expected a tensor to unravel, got .*\.'):
+        unravel_func(1)
+
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected a tensor of shape .*, got .*\.',
+    ):
+        unravel_func(flat.reshape((-1, 1)))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected a tensor of shape .*, got .*\.',
+    ):
+        unravel_func(torch.cat([flat, torch.zeros((1,))]))
+
+    if all(leaf.dtype == flat.dtype for leaf in leaves):
+        unravel_func(flat.to(torch.complex128))
+    else:
+        with pytest.raises(
+            ValueError,
+            match=r'The unravel function expected a tensor of dtype .*, got dtype .*\.',
+        ):
+            unravel_func(flat.to(torch.complex128))
+
+
+@parametrize(tree=list(TREES + LEAVES))
+def test_tree_ravel_single_dtype(tree):
+    random.seed(0)
+    dtype = torch.float16
+    default_dtype = torch.tensor([]).dtype
+
+    def replace_leaf(_):
+        candidates = []
+        shapes = [
+            (),
+            (random.randint(1, 10),),
+            (random.randint(1, 10), random.randint(1, 10)),
+            (random.randint(1, 10), random.randint(1, 10), random.randint(1, 10)),
+        ]
+        candidates.extend(
+            (5.0 * (2.0 * torch.randn(size=shape) - 1.0)).to(dtype) for shape in shapes
+        )
+        return random.choice(candidates)
+
+    tree = optree.tree_map(replace_leaf, tree)
+    flat, unravel_func = optree.integration.torch.tree_ravel(tree)
+
+    leaves, treespec = optree.tree_flatten(tree)
+    assert flat.dtype == dtype if leaves else default_dtype
+    assert flat.numel() == sum(leaf.numel() for leaf in leaves)
+    assert flat.shape == (flat.numel(),)
+
+    reconstructed = unravel_func(flat)
+    reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed)
+    assert reconstructed_treespec == treespec
+    assert len(leaves) == len(reconstructed_leaves)
+    for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves):
+        assert torch.is_tensor(leaf)
+        assert torch.is_tensor(reconstructed_leaf)
+        assert torch.allclose(leaf, reconstructed_leaf)
+        assert leaf.dtype == reconstructed_leaf.dtype
+        assert leaf.shape == reconstructed_leaf.shape
+
+    with pytest.raises(ValueError, match=r'Expected a tensor to unravel, got .*\.'):
+        unravel_func(1)
+
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected a tensor of shape .*, got .*\.',
+    ):
+        unravel_func(flat.reshape((-1, 1)))
+    with pytest.raises(
+        ValueError,
+        match=r'The unravel function expected a tensor of shape .*, got .*\.',
+    ):
+        unravel_func(torch.cat([flat, torch.zeros((1,))]))
+
+    unravel_func(flat.to(torch.complex128))
+
+
+def test_tree_ravel_non_tensor():
+    with pytest.raises(ValueError, match=r'All leaves must be tensors\.'):
+        optree.integration.torch.tree_ravel(1)
+
+    with pytest.raises(ValueError, match=r'All leaves must be tensors\.'):
+        optree.integration.torch.tree_ravel((1, 2))
+
+    with pytest.raises(ValueError, match=r'All leaves must be tensors\.'):
+        optree.integration.torch.tree_ravel((torch.tensor(1), 2))
+
+    optree.integration.torch.tree_ravel((torch.tensor(1), torch.tensor(2)))