[Refactor] Update vmap imports to torch (#1102)

pytorch · Apr 28, 2023 · 4793640 · 4793640
1 parent 8d8256a
commit 4793640
Show file tree

Hide file tree

Showing 8 changed files with 26 additions and 8 deletions.
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
@@ -29,7 +29,10 @@
 
 _has_functorch = False
 try:
-    from torch import vmap
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
 
     _has_functorch = True
 except ImportError:

diff --git a/torchrl/modules/tensordict_module/common.py b/torchrl/modules/tensordict_module/common.py
@@ -180,7 +180,7 @@ class SafeModule(TensorDictModule):
     One can use a vmap operator to call the functional module. In this case the tensordict is expanded to match the
     batch size (i.e. the tensordict isn't modified in-place anymore):
         >>> # Model ensemble using vmap
-        >>> from functorch import vmap
+        >>> from torch import vmap
         >>> params_repeat = params.expand(4, *params.shape)
         >>> td_vmap = vmap(td_fmodule, (None, 0))(td.clone(), params_repeat)
         >>> print(td_vmap)

diff --git a/torchrl/modules/tensordict_module/sequence.py b/torchrl/modules/tensordict_module/sequence.py
@@ -83,7 +83,7 @@ class SafeSequential(TensorDictSequential, SafeModule):
                 shape=torch.Size([8]), space=None, device=cpu, dtype=torch.float32, domain=continuous))
 
     In the vmap case:
-        >>> from functorch import vmap
+        >>> from torch import vmap
         >>> params = params.expand(4, *params.shape)
         >>> td_vmap = vmap(td_module, (None, 0))(td, params)
         >>> print(td_vmap)

diff --git a/torchrl/objectives/deprecated.py b/torchrl/objectives/deprecated.py
@@ -26,7 +26,10 @@
 from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
 
 try:
-    from functorch import vmap
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
 
     FUNCTORCH_ERR = ""
     _has_functorch = True

diff --git a/torchrl/objectives/iql.py b/torchrl/objectives/iql.py
@@ -23,7 +23,10 @@
 from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
 
 try:
-    from functorch import vmap
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
 
     _has_functorch = True
     err = ""

diff --git a/torchrl/objectives/redq.py b/torchrl/objectives/redq.py
@@ -25,7 +25,10 @@
 from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
 
 try:
-    from functorch import vmap
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
 
     FUNCTORCH_ERR = ""
     _has_functorch = True

diff --git a/torchrl/objectives/sac.py b/torchrl/objectives/sac.py
@@ -27,7 +27,10 @@
 from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
 
 try:
-    from functorch import vmap
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
 
     _has_functorch = True
     err = ""

diff --git a/torchrl/objectives/td3.py b/torchrl/objectives/td3.py
@@ -20,7 +20,10 @@
 from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
 
 try:
-    from functorch import vmap
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
 
     FUNCTORCH_ERR = ""
     _has_functorch = True