Changes due to deprecations in Cubed

cubed-dev · Oct 29, 2024 · d311ff6 · d311ff6
1 parent 0e9a6f3
commit d311ff6
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 9 deletions.
diff --git a/tests/benchmarks/test_array.py b/tests/benchmarks/test_array.py
@@ -8,7 +8,7 @@
 import cubed
 import cubed.random
 from cubed.core.optimization import multiple_inputs_optimize_dag, simple_optimize_dag
-from cubed.extensions.rich import RichProgressBar
+from cubed.diagnostics.rich import RichProgressBar
 
 from ..utils import run
 
@@ -47,9 +47,7 @@ def test_quadratic_means_xarray(tmp_path, runtime, benchmark_all, optimizer, t_l
     quad = ds**2
     quad["uv"] = ds.anom_u * ds.anom_v
     print(quad)
-    result = quad.mean(
-        "time", skipna=False, use_new_impl=True, split_every=10
-    )
+    result = quad.mean("time", skipna=False, split_every=10)
 
     if optimizer == "new-optimizer":
         opt_fn = partial(multiple_inputs_optimize_dag, max_total_num_input_blocks=20)

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -212,7 +212,7 @@ def benchmark_memory(test_run_benchmark):
     .. code-block:: python
 
         def test_something(benchmark_memory):
-            history = cubed.extensions.history.HistoryCallback()
+            history = cubed.diagnostics.history.HistoryCallback()
             with benchmark_memory(history):
                 cubed.compute(*arrs, callbacks=[history])
     """
@@ -259,7 +259,7 @@ def benchmark_tasks(test_run_benchmark):
     .. code-block:: python
 
         def test_something(benchmark_tasks):
-            history = cubed.extensions.history.HistoryCallback()
+            history = cubed.diagnostics.history.HistoryCallback()
             with benchmark_tasks(history):
                 cubed.compute(*arrs, callbacks=[history])
     """
@@ -305,7 +305,7 @@ def benchmark_all(
     .. code-block:: python
 
         def test_something(benchmark_all):
-            history = cubed.extensions.history.HistoryCallback()
+            history = cubed.diagnostics.history.HistoryCallback()
             with benchmark_all(history):
                 cubed.compute(*arrs, callbacks=[history])
     

diff --git a/tests/utils.py b/tests/utils.py
@@ -2,8 +2,8 @@
 import yaml
 
 import cubed
-from cubed.extensions.history import HistoryCallback
-from cubed.extensions.timeline import TimelineVisualizationCallback
+from cubed.diagnostics.history import HistoryCallback
+from cubed.diagnostics.timeline import TimelineVisualizationCallback
 from cubed.spec import spec_from_config
 from cubed import config