add torch "expand_dims"

jcmgray · Aug 1, 2023 · 9f3065a · 9f3065a
1 parent 84b070c
commit 9f3065a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/autoray/autoray.py b/autoray/autoray.py
@@ -2003,6 +2003,7 @@ def numpy_like(N, M=None, dtype=None, **kwargs):
 _FUNC_ALIASES["torch", "linalg.expm"] = "matrix_exp"
 _FUNC_ALIASES["torch", "conjugate"] = "conj"
 _FUNC_ALIASES["torch", "split"] = "tensor_split"
+_FUNC_ALIASES["torch", "expand_dims"] = "unsqueeze"
 
 _SUBMODULE_ALIASES["torch", "linalg.expm"] = "torch"
 _SUBMODULE_ALIASES["torch", "random.normal"] = "torch"
@@ -2034,6 +2035,9 @@ def numpy_like(N, M=None, dtype=None, **kwargs):
 _CUSTOM_WRAPPERS["torch", "take"] = make_translator(
     [("a", ("input",)), ("indices", ("index",)), ("axis", ("dim",))]
 )
+_CUSTOM_WRAPPERS["torch", "expand_dims"] = make_translator(
+    [("a", ("input",)), ("axis", ("dim",))]
+)
 
 # for older versions of torch, can provide some alternative implementations
 _MODULE_ALIASES["torch[alt]"] = "torch"