dev(narugo): add check for torch 2

.github/workflows/doc.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ 3.7 ]
+        python-version: [ 3.8 ]
 
     services:
       plantuml:

test/torch/funcs/test_wrapper.py

@@ -98,3 +98,9 @@ def test_vmap_nested(self, treetensor_x, treetensor_y):
                 }
             })
         ).all()
+
+    @skipUnless(vpip('torch') < '2', 'Torch 1.x required.')
+    def test_vmap_torch_1x(self, treetensor_x, treetensor_y):
+        f = lambda x, y: (x.sum() + y.mean() * 2)
+        with pytest.raises(NotImplementedError):
+            _ = ttorch.vmap(f)

treetensor/torch/funcs/wrapper.py

@@ -1,13 +1,19 @@
 import torch
+from hbutils.testing import vpip
 
 from .base import doc_from_base, wrap_for_treelize
 
 __all__ = [
     'vmap',
 ]
 
+_is_torch_2 = vpip('torch') >= '2'
+
 
 @doc_from_base()
 @wrap_for_treelize()
 def vmap(func, *args, **kwargs):
-    return torch.vmap(func, *args, **kwargs)
+    if _is_torch_2:
+        return torch.vmap(func, *args, **kwargs)
+    else:
+        raise NotImplementedError(f'Function vmap is not supported in torch {torch.__version__}.')