dev(narugo): fix issues

opendilab · Sep 19, 2023 · 114f98b · 114f98b
1 parent f0ffabe
commit 114f98b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 17 deletions.
diff --git a/test/torch/funcs/test_construct.py b/test/torch/funcs/test_construct.py
@@ -219,9 +219,9 @@ def test_rand_like(self):
 
         _target = ttorch.rand_like({
             'a': torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32),
-            'b': torch.tensor([1, 2, 3, 4], dtype=torch.float32),
+            'b': torch.tensor([1, 2, 3, 4], dtype=torch.float),
             'x': {
-                'c': torch.tensor([5, 6, 7], dtype=torch.float32),
+                'c': torch.tensor([5, 6, 7], dtype=torch.float64),
                 'd': torch.tensor([[[8, 9]]], dtype=torch.float32),
             }
         })

diff --git a/test/torch/funcs/test_wrapper.py b/test/torch/funcs/test_wrapper.py
@@ -33,9 +33,9 @@ class TestTorchTensorWrapper:
     @skipUnless(vpip('torch') >= '2', 'Torch 2 required.')
     def test_vmap(self, treetensor_x, treetensor_y):
         f = lambda x, y: (x.sum() + y.mean() * 2)
-        n_pow = torch.vmap(f)
-        batched_pow = ttorch.vmap(f)
-        r = batched_pow(treetensor_x, treetensor_y)
+        native_vf = torch.vmap(f)
+        tv_vf = ttorch.vmap(f)
+        r = tv_vf(treetensor_x, treetensor_y)
 
         assert r.shape == Size({
             'a': (2,),
@@ -46,19 +46,19 @@ def test_vmap(self, treetensor_x, treetensor_y):
         assert ttorch.isclose(
             r,
             ttorch.tensor({
-                'a': n_pow(treetensor_x.a, treetensor_y.a),
+                'a': native_vf(treetensor_x.a, treetensor_y.a),
                 'b': {
-                    'x': n_pow(treetensor_x.b.x, treetensor_y.b.x),
+                    'x': native_vf(treetensor_x.b.x, treetensor_y.b.x),
                 }
             })
         ).all()
 
     @skipUnless(vpip('torch') >= '2', 'Torch 2 required.')
     def test_vmap_in_dims(self, treetensor_x, treetensor_y):
         f = lambda x, y: (x.sum() + y.mean() * 2)
-        n_pow = torch.vmap(f, in_dims=1)
-        batched_pow = ttorch.vmap(f, in_dims=1)
-        r = batched_pow(treetensor_x, treetensor_y)
+        native_vf = torch.vmap(f, in_dims=1)
+        tv_vf = ttorch.vmap(f, in_dims=1)
+        r = tv_vf(treetensor_x, treetensor_y)
 
         assert r.shape == Size({
             'a': (5,),
@@ -69,19 +69,19 @@ def test_vmap_in_dims(self, treetensor_x, treetensor_y):
         assert ttorch.isclose(
             r,
             ttorch.tensor({
-                'a': n_pow(treetensor_x.a, treetensor_y.a),
+                'a': native_vf(treetensor_x.a, treetensor_y.a),
                 'b': {
-                    'x': n_pow(treetensor_x.b.x, treetensor_y.b.x),
+                    'x': native_vf(treetensor_x.b.x, treetensor_y.b.x),
                 }
             })
         ).all()
 
     @skipUnless(vpip('torch') >= '2', 'Torch 2 required.')
     def test_vmap_nested(self, treetensor_x, treetensor_y):
         f = lambda x, y: (x.sum() + y.mean() * 2)
-        n_pow = torch.vmap(torch.vmap(f))
-        batched_pow = ttorch.vmap(ttorch.vmap(f))
-        r = batched_pow(treetensor_x, treetensor_y)
+        native_vf = torch.vmap(torch.vmap(f))
+        tv_vf = ttorch.vmap(ttorch.vmap(f))
+        r = tv_vf(treetensor_x, treetensor_y)
 
         assert r.shape == Size({
             'a': (2, 5),
@@ -92,9 +92,9 @@ def test_vmap_nested(self, treetensor_x, treetensor_y):
         assert ttorch.isclose(
             r,
             ttorch.tensor({
-                'a': n_pow(treetensor_x.a, treetensor_y.a),
+                'a': native_vf(treetensor_x.a, treetensor_y.a),
                 'b': {
-                    'x': n_pow(treetensor_x.b.x, treetensor_y.b.x),
+                    'x': native_vf(treetensor_x.b.x, treetensor_y.b.x),
                 }
             })
         ).all()