replace 'cache_mode' option in tests with 'cache' (#2744)

Author: beverlylytle

thunder/tests/distributed/test_ddp.py

@@ -165,7 +165,7 @@ def get_model_and_optimizer(device):
             m = ToyModel().to(device)
             jitted_m = thunder.jit(
                 m,
-                cache_mode=CACHE_OPTIONS.CONSTANT_VALUES,
+                cache=CACHE_OPTIONS.CONSTANT_VALUES,
                 executors=executors_map[executor].executors_list(),
             )
             jitted_ddp_m = ddp(jitted_m, bucket_size_in_mb=bucket_size_in_mb)

thunder/tests/distributed/test_fsdp.py

@@ -137,7 +137,7 @@ def get_model_and_optimizer(device):
             m = ToyModel().to(device)
             jitted_m = thunder.jit(
                 m,
-                cache_mode=CACHE_OPTIONS.CONSTANT_VALUES,
+                cache=CACHE_OPTIONS.CONSTANT_VALUES,
                 executors=executors_map[executor].executors_list(),
             )
             jitted_fsdp_m = fsdp(jitted_m, bucketing_strategy=bucketing_strategy, sharding_strategy=fsdptype)

thunder/tests/test_core.py

@@ -827,7 +827,7 @@ def test_static_caching(executor, device: str, dtype: dtypes.dtype):
     def foo(a, b):
         return a + b
 
-    cfoo = thunder.jit(foo, cache_mode="constant values")
+    cfoo = thunder.jit(foo, cache="constant values")
 
     assert cache_option(cfoo) == thunder.CACHE_OPTIONS.CONSTANT_VALUES
 
@@ -894,7 +894,7 @@ def foo(a, b):
     def bar(a, b):
         return a, b
 
-    cbar = thunder.jit(bar, cache_mode="constant values")
+    cbar = thunder.jit(bar, cache="constant values")
 
     astr = "a"
     bstr = "b"
@@ -928,7 +928,7 @@ def bar(a, b):
 
     # Module tests
     m = torch.nn.Linear(5, 5, device=device, dtype=torch_dtype)
-    cm = thunder.jit(m, cache_mode="constant values")
+    cm = thunder.jit(m, cache="constant values")
 
     inp = make_tensor((5, 5), device=device, dtype=torch_dtype)
 
@@ -979,7 +979,7 @@ def caz(tup):
             accum += x
         return accum
 
-    ccaz = thunder.jit(caz, cache_mode="constant values")
+    ccaz = thunder.jit(caz, cache="constant values")
 
     inp0 = [5, 3, 7]
     thunder_result = ccaz(inp0)
@@ -1026,7 +1026,7 @@ def caz(tup):
     def daz(*, a, b):
         return a + b
 
-    cdaz = thunder.jit(daz, cache_mode="constant values")
+    cdaz = thunder.jit(daz, cache="constant values")
 
     inp0 = {"a": a, "b": b}
     thunder_result = cdaz(**inp0)

thunder/tests/test_grad.py

@@ -1336,11 +1336,11 @@ def test_torch_autograd_module(executor, device, _):
     a = make_tensor((2, 3), device=device, dtype=torch.float32, requires_grad=True)
     g = make_tensor((2, 4), device=device, dtype=torch.float32)
 
-    for cache_mode in ("constant values", "same input"):
+    for cache_mode in ("constant values", "no caching"):
         lc = executor.make_callable(
             linear,
             disable_torch_autograd=False,
-            cache_mode=cache_mode,
+            cache=cache_mode,
         )
         lc.zero_grad()
         a.grad = None