Switch to torch.compile being the default for the python backend

Author: ar4

src/deepwave/elastic.py

@@ -442,8 +442,8 @@ def elastic(
             Can be a string specifying whether to use PyTorch's JIT ("jit"),
             torch.compile ("compile"), or eager mode ("eager"). Alternatively
             a boolean can be provided, with True using the Python backend
-            with JIT, while the default, False, instead uses the compiled
-            C/CUDA.
+            with torch.compile, while the default, False, instead uses the
+            compiled C/CUDA.
 
     Returns:
         Tuple:
@@ -2317,9 +2317,9 @@ def elastic_python(
     receivers_p_i_masked = torch.zeros_like(receivers_p_i)
     receivers_p_i_masked[receiver_p_mask] = receivers_p_i[receiver_p_mask]
 
-    rdy = torch.tensor(1 / dy, dtype=dtype)
-    rdx = torch.tensor(1 / dx, dtype=dtype)
-    dt_tensor = torch.tensor(dt, dtype=dtype)
+    rdy = torch.tensor(1 / dy, dtype=dtype, device=device)
+    rdx = torch.tensor(1 / dx, dtype=dtype, device=device)
+    dt_tensor = torch.tensor(dt, dtype=dtype, device=device)
 
     for step in range(nt // step_ratio):
         if forward_callback is not None and step % callback_frequency == 0:
@@ -2522,11 +2522,9 @@ def elastic_func(
     global _update_velocities_jit, _update_velocities_compile, _update_velocities_opt
     global _update_stresses_jit, _update_stresses_compile, _update_stresses_opt
 
-    if not python_backend:
-        func = ElasticForwardFunc.apply
-    else:
+    if python_backend:
         if python_backend is True:
-            mode = "jit"
+            mode = "compile"
         elif isinstance(python_backend, str):
             mode = python_backend.lower()
         else:
@@ -2558,7 +2556,7 @@ def elastic_func(
         else:
             raise ValueError(f"Unknown python_backend value {mode!r}.")
 
-        func = elastic_python
+    func = elastic_python if python_backend else ElasticForwardFunc.apply
 
     return cast(
         "Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, "

src/deepwave/regular_grid.py

@@ -4,7 +4,7 @@
 profiles for wave propagation simulations on a regular grid.
 """
 
-from typing import List, Union
+from typing import List
 
 import torch
 
@@ -83,8 +83,12 @@ def set_pml_profiles(
         device,
         pml_freq,
     )
-    dbydy = diffx1(by, accuracy, 1 / grid_spacing[0])
-    dbxdx = diffx1(bx, accuracy, 1 / grid_spacing[1])
+    dbydy = diffx1(
+        by, accuracy, torch.tensor(1 / grid_spacing[0], dtype=dtype, device=device)
+    )
+    dbxdx = diffx1(
+        bx, accuracy, torch.tensor(1 / grid_spacing[1], dtype=dtype, device=device)
+    )
     ay = ay[None, :, None]
     ax = ax[None, None, :]
     by = by[None, :, None]
@@ -94,9 +98,7 @@ def set_pml_profiles(
     return [ay, ax, by, bx, dbydy, dbxdx]
 
 
-def diffy1(
-    a: torch.Tensor, accuracy: int, rdy: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffy1(a: torch.Tensor, accuracy: int, rdy: torch.Tensor) -> torch.Tensor:
     """Calculates the first derivative in the y-direction."""
     if accuracy == 2:
         return torch.nn.functional.pad(
@@ -133,9 +135,7 @@ def diffy1(
     )
 
 
-def diffx1(
-    a: torch.Tensor, accuracy: int, rdx: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffx1(a: torch.Tensor, accuracy: int, rdx: torch.Tensor) -> torch.Tensor:
     """Calculates the first derivative in the x-direction."""
     if accuracy == 2:
         return torch.nn.functional.pad(
@@ -172,9 +172,7 @@ def diffx1(
     )
 
 
-def diffy2(
-    a: torch.Tensor, accuracy: int, rdy2: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffy2(a: torch.Tensor, accuracy: int, rdy2: torch.Tensor) -> torch.Tensor:
     """Calculates the second derivative in the y-direction."""
     if accuracy == 2:
         return torch.nn.functional.pad(
@@ -215,9 +213,7 @@ def diffy2(
     )
 
 
-def diffx2(
-    a: torch.Tensor, accuracy: int, rdx2: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffx2(a: torch.Tensor, accuracy: int, rdx2: torch.Tensor) -> torch.Tensor:
     """Calculates the second derivative in the x-direction."""
     if accuracy == 2:
         return torch.nn.functional.pad(

src/deepwave/scalar.py

@@ -366,8 +366,8 @@ def scalar(
             Can be a string specifying whether to use PyTorch's JIT ("jit"),
             torch.compile ("compile"), or eager mode ("eager"). Alternatively
             a boolean can be provided, with True using the Python backend
-            with JIT, while the default, False, instead uses the compiled
-            C/CUDA.
+            with torch.compile, while the default, False, instead uses the
+            compiled C/CUDA.
 
     Returns:
         Tuple:
@@ -2233,11 +2233,11 @@ def scalar_python(
     nx = v.shape[-1]
     n_receivers_per_shot = receivers_i.numel() // n_shots
     receiver_amplitudes = torch.empty(0, device=device, dtype=dtype)
-    rdy = torch.tensor(1 / dy, dtype=dtype)
-    rdx = torch.tensor(1 / dx, dtype=dtype)
+    rdy = torch.tensor(1 / dy, dtype=dtype, device=device)
+    rdx = torch.tensor(1 / dx, dtype=dtype, device=device)
     rdy2 = rdy**2
     rdx2 = rdx**2
-    dt_tensor = torch.tensor(dt, dtype=dtype)
+    dt_tensor = torch.tensor(dt, dtype=dtype, device=device)
 
     if receivers_i.numel() > 0:
         receiver_amplitudes.resize_(nt, n_shots, n_receivers_per_shot)
@@ -2346,11 +2346,9 @@ def scalar_func(
     """
     global _forward_step_jit, _forward_step_compile, _forward_step_opt
 
-    if not python_backend:
-        func = ScalarForwardFunc.apply
-    else:
+    if python_backend:
         if python_backend is True:
-            mode = "jit"
+            mode = "compile"
         elif isinstance(python_backend, str):
             mode = python_backend.lower()
         else:
@@ -2371,7 +2369,7 @@ def scalar_func(
         else:
             raise ValueError(f"Unknown python_backend value {mode!r}.")
 
-        func = scalar_python
+    func = scalar_python if python_backend else ScalarForwardFunc.apply
 
     return cast(
         "Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, "

src/deepwave/scalar_born.py

@@ -270,8 +270,8 @@ def scalar_born(
             Can be a string specifying whether to use PyTorch's JIT ("jit"),
             torch.compile ("compile"), or eager mode ("eager"). Alternatively
             a boolean can be provided, with True using the Python backend
-            with JIT, while the default, False, instead uses the compiled
-            C/CUDA.
+            with torch.compile, while the default, False, instead uses the
+            compiled C/CUDA.
 
     Returns:
         Tuple:
@@ -1795,11 +1795,11 @@ def scalar_born_python(
     n_receiverssc_per_shot = receiverssc_i.numel() // n_shots
     receiver_amplitudes = torch.empty(0, device=device, dtype=dtype)
     receiver_amplitudessc = torch.empty(0, device=device, dtype=dtype)
-    rdy = torch.tensor(1 / dy, dtype=dtype)
-    rdx = torch.tensor(1 / dx, dtype=dtype)
+    rdy = torch.tensor(1 / dy, dtype=dtype, device=device)
+    rdx = torch.tensor(1 / dx, dtype=dtype, device=device)
     rdy2 = rdy**2
     rdx2 = rdx**2
-    dt_tensor = torch.tensor(dt, dtype=dtype)
+    dt_tensor = torch.tensor(dt, dtype=dtype, device=device)
 
     if receivers_i.numel() > 0:
         receiver_amplitudes.resize_(nt, n_shots, n_receivers_per_shot)
@@ -1977,11 +1977,9 @@ def scalar_born_func(
     """
     global _forward_step_jit, _forward_step_compile, _forward_step_opt
 
-    if not python_backend:
-        func = ScalarBornForwardFunc.apply
-    else:
+    if python_backend:
         if python_backend is True:
-            mode = "jit"
+            mode = "compile"
         elif isinstance(python_backend, str):
             mode = python_backend.lower()
         else:
@@ -2002,7 +2000,7 @@ def scalar_born_func(
         else:
             raise ValueError(f"Unknown python_backend value {mode!r}.")
 
-        func = scalar_born_python
+    func = scalar_born_python if python_backend else ScalarBornForwardFunc.apply
 
     return cast(
         "Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, "

src/deepwave/staggered_grid.py

@@ -4,7 +4,7 @@
 profiles for wave propagation simulations on a staggered grid.
 """
 
-from typing import List, Union
+from typing import List
 
 import torch
 
@@ -122,9 +122,7 @@ def set_pml_profiles(
     return [ay, ayh, ax, axh, by, byh, bx, bxh]
 
 
-def diffy1(
-    a: torch.Tensor, accuracy: int, rdy: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffy1(a: torch.Tensor, accuracy: int, rdy: torch.Tensor) -> torch.Tensor:
     """Calculates the first y derivative at integer grid points."""
     if accuracy == 2:
         return torch.nn.functional.pad(
@@ -161,9 +159,7 @@ def diffy1(
     )
 
 
-def diffx1(
-    a: torch.Tensor, accuracy: int, rdx: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffx1(a: torch.Tensor, accuracy: int, rdx: torch.Tensor) -> torch.Tensor:
     """Calculates the first x derivative at integer grid points."""
     if accuracy == 2:
         return torch.nn.functional.pad((a[..., 1:] - a[..., :-1]) * rdx, (1, 0))
@@ -198,9 +194,7 @@ def diffx1(
     )
 
 
-def diffyh1(
-    a: torch.Tensor, accuracy: int, rdy: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffyh1(a: torch.Tensor, accuracy: int, rdy: torch.Tensor) -> torch.Tensor:
     """Calculates the first y derivative at half integer grid points."""
     if accuracy == 2:
         return torch.nn.functional.pad(
@@ -237,9 +231,7 @@ def diffyh1(
     )
 
 
-def diffxh1(
-    a: torch.Tensor, accuracy: int, rdx: Union[torch.Tensor, float]
-) -> torch.Tensor:
+def diffxh1(a: torch.Tensor, accuracy: int, rdx: torch.Tensor) -> torch.Tensor:
     """Calculates the first x derivative at half integer grid points."""
     if accuracy == 2:
         return torch.nn.functional.pad((a[..., 2:] - a[..., 1:-1]) * rdx, (1, 1))

tests/test_elastic.py

@@ -19,6 +19,8 @@
 )
 from deepwave.wavelets import ricker
 
+torch._dynamo.config.cache_size_limit = 256  # noqa: SLF001
+
 DEFAULT_LAMB = 550000000
 DEFAULT_MU = 2200000000
 DEFAULT_BUOYANCY = 1 / 2200

tests/test_scalar.py

@@ -10,6 +10,8 @@
 from deepwave import IGNORE_LOCATION, Scalar, scalar
 from deepwave.common import cfl_condition, downsample, upsample
 
+torch._dynamo.config.cache_size_limit = 256  # noqa: SLF001
+
 
 def scalarprop(
     model: torch.Tensor,
@@ -536,7 +538,7 @@ def wrap(python):
     torch.manual_seed(1)
     out_compiled = wrap(False)
     torch.manual_seed(1)
-    out_python = wrap(True)
+    out_python = wrap("jit")
     for oc, op in zip(out_compiled, out_python):
         assert torch.allclose(oc, op)
 
@@ -1286,7 +1288,7 @@ def wrap(python):
         return grads + grad_grads
 
     torch.manual_seed(1)
-    out_python = wrap(True)
+    out_python = wrap("jit")
     torch.manual_seed(1)
     out_compiled = wrap(False)
     for oc, op in zip(out_compiled, out_python):

tests/test_scalar_born.py

@@ -8,6 +8,8 @@
 from deepwave import IGNORE_LOCATION, ScalarBorn, scalar_born
 from deepwave.common import cfl_condition, downsample, upsample
 
+torch._dynamo.config.cache_size_limit = 256  # noqa: SLF001
+
 
 def test_born_scatter_2d() -> None:
     """Test Born propagation in a 2D model with a point scatterer."""