add scheduler classes

src/finch/__init__.py

@@ -113,6 +113,8 @@
     compiled,
     compute,
     set_optimizer,
+    DefaultScheduler,
+    GalleyScheduler,
 )
 from .dtypes import (
     int_,

src/finch/compiled.py

@@ -4,12 +4,6 @@
 from .tensor import Tensor
 
 
-def get_scheduler(name, verbose=False):
-    if name == "default":
-        return jl.Finch.default_scheduler()
-    elif name == "galley":
-        return jl.Finch.galley_scheduler(verbose=verbose)
-
 def compiled(opt=""):
     def inner(func):  
         @wraps(func)  
@@ -33,17 +27,31 @@ def lazy(tensor: Tensor):
         return Tensor(jl.Finch.LazyTensor(tensor._obj))
     return tensor
 
-def set_optimizer(opt="default"):
-    if opt == "default":
-        jl.Finch.set_scheduler_b(jl.Finch.default_scheduler())
-    elif opt == "galley":
-        jl.Finch.set_scheduler_b(jl.Finch.galley_scheduler())
+class AbstractScheduler():
+    pass
+
+class GalleyScheduler(AbstractScheduler):
+    def __init__(self, verbose=False):
+        self.verbose=verbose
+
+class DefaultScheduler(AbstractScheduler):
+    def __init__(self, verbose=False):
+        self.verbose=verbose
+
+def get_julia_scheduler(opt):
+    if isinstance(opt, DefaultScheduler):
+        return jl.Finch.default_scheduler(verbose=opt.verbose)
+    elif isinstance(opt, GalleyScheduler):
+        return jl.Finch.galley_scheduler(verbose=opt.verbose)
+
+def set_optimizer(opt):
+    jl.Finch.set_scheduler_b(get_julia_scheduler(opt))
     return
 
-def compute(tensor: Tensor, *, verbose: bool = False, opt="", tag=-1):
+def compute(tensor: Tensor, *, verbose: bool = False, opt=None, tag=-1):
     if not tensor.is_computed():
-        if opt == "":
+        if opt == None:
             return Tensor(jl.Finch.compute(tensor._obj, verbose=verbose, tag=tag))
         else:            
-            return Tensor(jl.Finch.compute(tensor._obj, verbose=verbose, tag=tag, ctx=get_scheduler(opt, verbose=verbose)))
+            return Tensor(jl.Finch.compute(tensor._obj, verbose=verbose, tag=tag, ctx=get_julia_scheduler(opt)))
     return tensor

tests/test_ops.py

@@ -9,7 +9,7 @@
 
 arr1d = np.array([1, 1, 2, 3])
 
-parametrize_optimizer = pytest.mark.parametrize("opt", ["default", "galley"])  
+parametrize_optimizer = pytest.mark.parametrize("opt", [finch.DefaultScheduler(), finch.GalleyScheduler()])  
 
 @parametrize_optimizer
 def test_eager(arr3d, opt):

tests/test_sparse.py

@@ -5,7 +5,7 @@
 
 import finch
 
-parametrize_optimizer =  pytest.mark.parametrize("opt", ["default", "galley"])
+parametrize_optimizer = pytest.mark.parametrize("opt", [finch.DefaultScheduler(), finch.GalleyScheduler()])  
 
 @pytest.mark.parametrize(
     "dtype,jl_dtype",