From 6aef851669a89cdb4dd694fe67a96f9d0155d31d Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 30 Apr 2026 13:49:23 -0400 Subject: [PATCH 1/6] fix: circular imports --- src/finchlite/__init__.py | 110 +++++++++++--------- src/finchlite/compile/bufferized_ndarray.py | 3 +- src/finchlite/interface/eager.py | 7 +- 3 files changed, 65 insertions(+), 55 deletions(-) diff --git a/src/finchlite/__init__.py b/src/finchlite/__init__.py index f050f8fb4..b7d3fd65b 100644 --- a/src/finchlite/__init__.py +++ b/src/finchlite/__init__.py @@ -1,55 +1,6 @@ -from .algebra import ( - FTyped, - Tensor, - TensorFType, - bool, - bool_, - complex64, - complex128, - complex_, - ffuncs, - finfo, - fisinstance, - float16, - float32, - float64, - float_, - ftype, - iinfo, - int8, - int16, - int32, - int64, - int_, - intp, - isdtype, - none_, - str_, - uint8, - uint16, - uint32, - uint64, -) -from .autoschedule.tensor_stats import ( - DC, - BlockedStats, - DCStats, - DenseStats, - TensorDef, - UniformStats, -) -from .codegen import ( - NumpyBuffer, - NumpyBufferFType, -) -from .compile import ( - AssemblyContext, - BufferizedNDArray, - Extent, - ExtentFType, - dimension, -) -from .finch_fused import jit +# isort: split +# interface must be imported first so EagerTensor is defined before the +# autoschedule→compile chain runs (compile.BufferizedNDArray inherits from it). from .interface import ( EagerTensor, LazyTensor, @@ -161,6 +112,61 @@ var, vecdot, ) + +# isort: split + +from .algebra import ( + FTyped, + Tensor, + TensorFType, + bool, + bool_, + complex64, + complex128, + complex_, + ffuncs, + finfo, + fisinstance, + float16, + float32, + float64, + float_, + ftype, + iinfo, + int8, + int16, + int32, + int64, + int_, + intp, + isdtype, + none_, + str_, + uint8, + uint16, + uint32, + uint64, +) +from .autoschedule.tensor_stats import ( + DC, + BlockedStats, + DCStats, + DenseStats, + TensorDef, + UniformStats, +) +from .codegen import ( + NumpyBuffer, + NumpyBufferFType, +) +from .compile import ( + AssemblyContext, + BufferizedNDArray, + Extent, + ExtentFType, + dimension, +) +from .finch_fused import jit from .tensor import ( DenseLevel, DenseLevelFType, diff --git a/src/finchlite/compile/bufferized_ndarray.py b/src/finchlite/compile/bufferized_ndarray.py index 9a8c00812..ba0f454e8 100644 --- a/src/finchlite/compile/bufferized_ndarray.py +++ b/src/finchlite/compile/bufferized_ndarray.py @@ -8,6 +8,7 @@ from ..algebra import FType, ImmutableStructFType, Tensor, TupleFType, ffuncs, ftype from ..codegen import NumpyBuffer, NumpyBufferFType from ..codegen.numba_codegen import to_numpy_type +from ..interface.eager import EagerTensor from . import looplets as lplt from .lower import AssemblyContext, FinchTensorFType @@ -16,7 +17,7 @@ def _get_default_strides(size: tuple[int, ...]) -> tuple[int, ...]: return tuple(np.cumprod((1,) + size[::-1]).astype(int))[-2::-1] -class BufferizedNDArray(Tensor): +class BufferizedNDArray(EagerTensor): def __init__( self, val: NumpyBuffer, diff --git a/src/finchlite/interface/eager.py b/src/finchlite/interface/eager.py index 70b6cec74..cca5b107c 100644 --- a/src/finchlite/interface/eager.py +++ b/src/finchlite/interface/eager.py @@ -6,8 +6,6 @@ from typing import Any from ..algebra import FinchOperator -from . import lazy -from .fuse import compute from .overrides import OverrideTensor @@ -231,6 +229,10 @@ def __ne__(self, other): return not_equal(self, other) +from . import lazy # noqa: E402 +from .fuse import compute # noqa: E402 + + def full( shape: int | tuple[int, ...], fill_value: bool | complex, @@ -418,6 +420,7 @@ def matmul(x1, x2, /): """ if isinstance(x1, lazy.LazyTensor) or isinstance(x2, lazy.LazyTensor): return lazy.matmul(x1, x2) + c = lazy.matmul(x1, x2) return compute(c) From 9f8e9b5efb3caf84e9c98ad1e55ea410bb2b4f5a Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 30 Apr 2026 14:11:21 -0400 Subject: [PATCH 2/6] fix: count_non_zeros --- src/finchlite/compile/bufferized_ndarray.py | 3 +++ src/finchlite/interface/eager.py | 10 +++++++++ tests/test_interface.py | 25 +++++++++++++++++++++ 3 files changed, 38 insertions(+) diff --git a/src/finchlite/compile/bufferized_ndarray.py b/src/finchlite/compile/bufferized_ndarray.py index ba0f454e8..419d46624 100644 --- a/src/finchlite/compile/bufferized_ndarray.py +++ b/src/finchlite/compile/bufferized_ndarray.py @@ -116,6 +116,9 @@ def __setitem__(self, index, value): index = np.ravel_multi_index(index, self._shape) self.val.store(index, value) + def count_nonzero(self, axis=None, *, keepdims=False): + return np.count_nonzero(self.to_numpy(), axis=axis, keepdims=keepdims) + def __str__(self): return f"{self.ftype}(shape={self.shape})" diff --git a/src/finchlite/interface/eager.py b/src/finchlite/interface/eager.py index cca5b107c..1373fe069 100644 --- a/src/finchlite/interface/eager.py +++ b/src/finchlite/interface/eager.py @@ -1045,6 +1045,16 @@ def not_equal(x1, x2): return compute(lazy.not_equal(x1, x2)) +def count_nonzero(a, axis=None, *, keepdims=False): + if isinstance(a, lazy.LazyTensor): + raise NotImplementedError("count_nonzero is not implemented for LazyTensor") + if hasattr(a, "count_nonzero"): + return a.count_nonzero(axis=axis, keepdims=keepdims) + raise NotImplementedError( + f"count_nonzero is not implemented for {type(a).__name__}" + ) + + def mean(x, /, *, axis: int | tuple[int, ...] | None = None, keepdims: bool = False): if isinstance(x, lazy.LazyTensor): return lazy.mean(x, axis=axis, keepdims=keepdims) diff --git a/tests/test_interface.py b/tests/test_interface.py index aa7122ae9..e361a6c83 100644 --- a/tests/test_interface.py +++ b/tests/test_interface.py @@ -650,6 +650,31 @@ def test_matmul(a, b, a_wrap, b_wrap): finch_assert_allclose(result_with_np, expected) +@pytest.mark.parametrize( + "a, b", + [ + ( + np.array([[1.0, 2.0], [3.0, 4.0]]), + np.array([[5.0, 6.0], [7.0, 8.0]]), + ), + ( + np.arange(12, dtype=np.float64).reshape(3, 4), + np.arange(8, dtype=np.float64).reshape(4, 2), + ), + ], +) +def test_matmul_bufferized_ndarray(a, b): + ba = finchlite.asarray(a) + bb = finchlite.asarray(b) + expected = a @ b + + result = finchlite.matmul(ba, bb) + result_with_op = ba @ bb + + finch_assert_allclose(result, expected) + finch_assert_allclose(result_with_op, expected) + + @pytest.mark.usefixtures("interpreter_scheduler") # TODO: remove @pytest.mark.parametrize( "a", From cda1efc140e7a0cb7219ffb805dffb059038e5c9 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 30 Apr 2026 15:50:57 -0400 Subject: [PATCH 3/6] fix: imports --- src/finchlite/__init__.py | 110 +++++++++----------- src/finchlite/compile/bufferized_ndarray.py | 3 - src/finchlite/interface/eager.py | 10 -- src/finchlite/interface/lazy.py | 3 +- 4 files changed, 54 insertions(+), 72 deletions(-) diff --git a/src/finchlite/__init__.py b/src/finchlite/__init__.py index b7d3fd65b..f050f8fb4 100644 --- a/src/finchlite/__init__.py +++ b/src/finchlite/__init__.py @@ -1,6 +1,55 @@ -# isort: split -# interface must be imported first so EagerTensor is defined before the -# autoschedule→compile chain runs (compile.BufferizedNDArray inherits from it). +from .algebra import ( + FTyped, + Tensor, + TensorFType, + bool, + bool_, + complex64, + complex128, + complex_, + ffuncs, + finfo, + fisinstance, + float16, + float32, + float64, + float_, + ftype, + iinfo, + int8, + int16, + int32, + int64, + int_, + intp, + isdtype, + none_, + str_, + uint8, + uint16, + uint32, + uint64, +) +from .autoschedule.tensor_stats import ( + DC, + BlockedStats, + DCStats, + DenseStats, + TensorDef, + UniformStats, +) +from .codegen import ( + NumpyBuffer, + NumpyBufferFType, +) +from .compile import ( + AssemblyContext, + BufferizedNDArray, + Extent, + ExtentFType, + dimension, +) +from .finch_fused import jit from .interface import ( EagerTensor, LazyTensor, @@ -112,61 +161,6 @@ var, vecdot, ) - -# isort: split - -from .algebra import ( - FTyped, - Tensor, - TensorFType, - bool, - bool_, - complex64, - complex128, - complex_, - ffuncs, - finfo, - fisinstance, - float16, - float32, - float64, - float_, - ftype, - iinfo, - int8, - int16, - int32, - int64, - int_, - intp, - isdtype, - none_, - str_, - uint8, - uint16, - uint32, - uint64, -) -from .autoschedule.tensor_stats import ( - DC, - BlockedStats, - DCStats, - DenseStats, - TensorDef, - UniformStats, -) -from .codegen import ( - NumpyBuffer, - NumpyBufferFType, -) -from .compile import ( - AssemblyContext, - BufferizedNDArray, - Extent, - ExtentFType, - dimension, -) -from .finch_fused import jit from .tensor import ( DenseLevel, DenseLevelFType, diff --git a/src/finchlite/compile/bufferized_ndarray.py b/src/finchlite/compile/bufferized_ndarray.py index 419d46624..ba0f454e8 100644 --- a/src/finchlite/compile/bufferized_ndarray.py +++ b/src/finchlite/compile/bufferized_ndarray.py @@ -116,9 +116,6 @@ def __setitem__(self, index, value): index = np.ravel_multi_index(index, self._shape) self.val.store(index, value) - def count_nonzero(self, axis=None, *, keepdims=False): - return np.count_nonzero(self.to_numpy(), axis=axis, keepdims=keepdims) - def __str__(self): return f"{self.ftype}(shape={self.shape})" diff --git a/src/finchlite/interface/eager.py b/src/finchlite/interface/eager.py index 1373fe069..cca5b107c 100644 --- a/src/finchlite/interface/eager.py +++ b/src/finchlite/interface/eager.py @@ -1045,16 +1045,6 @@ def not_equal(x1, x2): return compute(lazy.not_equal(x1, x2)) -def count_nonzero(a, axis=None, *, keepdims=False): - if isinstance(a, lazy.LazyTensor): - raise NotImplementedError("count_nonzero is not implemented for LazyTensor") - if hasattr(a, "count_nonzero"): - return a.count_nonzero(axis=axis, keepdims=keepdims) - raise NotImplementedError( - f"count_nonzero is not implemented for {type(a).__name__}" - ) - - def mean(x, /, *, axis: int | tuple[int, ...] | None = None, keepdims: bool = False): if isinstance(x, lazy.LazyTensor): return lazy.mean(x, axis=axis, keepdims=keepdims) diff --git a/src/finchlite/interface/lazy.py b/src/finchlite/interface/lazy.py index 918eebee6..1c620d89c 100644 --- a/src/finchlite/interface/lazy.py +++ b/src/finchlite/interface/lazy.py @@ -28,7 +28,6 @@ ) from ..algebra.ftypes import FDType from ..autoschedule.tensor_stats import StatsInterpreter -from ..compile import BufferizedNDArray from ..finch_logic import ( Aggregate, Alias, @@ -463,6 +462,8 @@ def asarray(arg: Any, format: TensorFType | None = None) -> Any: from finchlite.interface.scalar import Scalar if isinstance(arg, np.ndarray): + from ..compile import BufferizedNDArray + return BufferizedNDArray.from_numpy(arg) if np.isscalar(arg) or arg is None: return Scalar(arg) From 39cbd5a794e57944fcffdafcb1b64917009bfa51 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 30 Apr 2026 16:27:12 -0400 Subject: [PATCH 4/6] fix: circular imports --- src/finchlite/autoschedule/compiler.py | 5 ++++- src/finchlite/autoschedule/formatter.py | 3 ++- src/finchlite/autoschedule/tensor_stats/dc_stats.py | 13 +++++++------ src/finchlite/compile/__init__.py | 4 +++- 4 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/finchlite/autoschedule/compiler.py b/src/finchlite/autoschedule/compiler.py index 07632a0a4..3ed146149 100644 --- a/src/finchlite/autoschedule/compiler.py +++ b/src/finchlite/autoschedule/compiler.py @@ -11,7 +11,6 @@ from .. import finch_logic as lgc from .. import finch_notation as ntn from ..algebra import FinchOperator, FType, ffuncs, ftypes -from ..compile.lower import make_extent from ..finch_assembly import AssemblyLibrary from ..finch_logic import LogicLoader, StatsFactory, TensorStats, compute_shape_vars from ..finch_notation import NotationInterpreter @@ -101,6 +100,8 @@ def _lower_query_of_reorder( arg: lgc.Table, reorder_idxs: tuple[lgc.Field, ...], ): + from ..compile.lower import make_extent + arg_dims = arg.dimmap(merge_shapes, self.shapes) shapes_map = dict(zip(arg.idxs, arg_dims, strict=True)) shapes = { @@ -189,6 +190,8 @@ def _lower_query_of_aggregate( agg_arg: lgc.Reorder, agg_idxs: tuple[lgc.Field, ...], ): + from ..compile.lower import make_extent + # Build a dict mapping fields to their shapes arg_dims = agg_arg.dimmap(merge_shapes, self.shapes) shapes_map = dict(zip(agg_arg.idxs, arg_dims, strict=True)) diff --git a/src/finchlite/autoschedule/formatter.py b/src/finchlite/autoschedule/formatter.py index f0cccf196..c19c45653 100644 --- a/src/finchlite/autoschedule/formatter.py +++ b/src/finchlite/autoschedule/formatter.py @@ -7,7 +7,6 @@ from .. import finch_logic as lgc from ..algebra import FType, TensorFType, TupleFType, ftype from ..codegen import NumpyBufferFType -from ..compile import BufferizedNDArrayFType from ..finch_assembly import AssemblyLibrary from ..finch_logic import LogicLoader, MockLogicLoader, StatsFactory, TensorStats from ..util.logging import LOG_LOGIC_POST_OPT @@ -95,6 +94,8 @@ def __init__( super().__init__(loader) def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[FType, ...]): + from ..compile.bufferized_ndarray import BufferizedNDArrayFType + return BufferizedNDArrayFType( buffer_type=NumpyBufferFType(ftype(fill_value)), ndim=len(shape_type), diff --git a/src/finchlite/autoschedule/tensor_stats/dc_stats.py b/src/finchlite/autoschedule/tensor_stats/dc_stats.py index 48c301473..b95f3d742 100644 --- a/src/finchlite/autoschedule/tensor_stats/dc_stats.py +++ b/src/finchlite/autoschedule/tensor_stats/dc_stats.py @@ -13,7 +13,6 @@ from ... import finch_notation as ntn from ...algebra import Tensor, ffuncs, ftype, int64 from ...algebra.algebra import FinchOperator -from ...compile import BufferizedNDArray, make_extent from .numeric_stats import NumericStats from .tensor_def import TensorDef from .tensor_stats import BaseTensorStatsFactory @@ -37,9 +36,6 @@ class DC: value: float -_INT64_VECTOR_FTYPE = BufferizedNDArray.from_numpy(np.zeros(1, dtype=np.int64)).ftype - - def _int_tuple_ftype(size: int): return ftype(tuple(np.int64(0) for _ in range(size))) @@ -152,17 +148,22 @@ def _structure_to_dcs(self, arr: Tensor, fields: Iterable[Field]) -> set[DC]: # For each field i, we compute DC({}, {i}) and DC({i}, {*fields}). # Additionally, we compute the nnz for the full tensor DC({}, {*fields}). def _array_to_dcs(self, arr: Any, fields: Iterable[Field]) -> set[DC]: + from ...compile import BufferizedNDArray, make_extent + + int64_vector_ftype = BufferizedNDArray.from_numpy( + np.zeros(1, dtype=np.int64) + ).ftype fields = list(fields) ndims = len(fields) dim_loop_variables = [ntn.Variable(f"{fields[i]}", int64) for i in range(ndims)] dim_array_variables = [ - ntn.Variable(f"x_{fields[i]}", _INT64_VECTOR_FTYPE) for i in range(ndims) + ntn.Variable(f"x_{fields[i]}", int64_vector_ftype) for i in range(ndims) ] dim_size_variables = [ ntn.Variable(f"n_{fields[i]}", int64) for i in range(ndims) ] dim_array_slots = [ - ntn.Slot(f"x_{fields[i]}_", _INT64_VECTOR_FTYPE) for i in range(ndims) + ntn.Slot(f"x_{fields[i]}_", int64_vector_ftype) for i in range(ndims) ] dim_proj_variables = [ ntn.Variable(f"proj_{fields[i]}", int64) for i in range(ndims) diff --git a/src/finchlite/compile/__init__.py b/src/finchlite/compile/__init__.py index b6c453bc4..c306a5652 100644 --- a/src/finchlite/compile/__init__.py +++ b/src/finchlite/compile/__init__.py @@ -1,4 +1,3 @@ -from .bufferized_ndarray import BufferizedNDArray, BufferizedNDArrayFType from .lower import ( AssemblyContext, Extent, @@ -9,6 +8,9 @@ make_extent, ) +# isort: split +from .bufferized_ndarray import BufferizedNDArray, BufferizedNDArrayFType + __all__ = [ "AssemblyContext", "BufferizedNDArray", From 8661ea4cc79a33b9502416bb484f6a4d26a1aea0 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 30 Apr 2026 16:33:31 -0400 Subject: [PATCH 5/6] fix: test --- tests/test_notation_interpreter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_notation_interpreter.py b/tests/test_notation_interpreter.py index 59b09bfb3..7deb40f6e 100644 --- a/tests/test_notation_interpreter.py +++ b/tests/test_notation_interpreter.py @@ -220,4 +220,4 @@ def test_count_nonfill_vector(a): mod = ntn.NotationInterpreter()(prgm) cnt = mod.count_nonfill_vector(a) - assert cnt == np.count_nonzero(a) + assert cnt == np.count_nonzero(a.to_numpy()) From d08ae18fe7d8c74808e943ef2f0b6ecadc82bb9f Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 14 May 2026 17:21:48 -0400 Subject: [PATCH 6/6] fix: refactor to avoid circular imports --- src/finchlite/__init__.py | 18 ++++++++++-------- src/finchlite/adv_autoschedulers/__init__.py | 0 .../tensor_stats/__init__.py | 14 +++++++------- .../tensor_stats/blocked_stats.py | 0 .../tensor_stats/database_stats.py | 0 .../tensor_stats/dc_stats.py | 3 +-- .../tensor_stats/dense_stat.py | 0 .../tensor_stats/dummy_stats.py | 0 .../tensor_stats/numeric_stats.py | 0 .../tensor_stats/stats_interpreter.py | 0 .../tensor_stats/tensor_def.py | 0 .../tensor_stats/tensor_stats.py | 0 .../tensor_stats/uniform_stats.py | 0 src/finchlite/autoschedule/cache.py | 2 +- src/finchlite/autoschedule/einsum.py | 3 +-- src/finchlite/autoschedule/executor.py | 3 ++- .../logical_optimizer/annotated_query.py | 2 +- src/finchlite/interface/lazy.py | 3 ++- tests/test_cache_loader.py | 9 ++++++--- tests/test_galley_branch_and_bound.py | 2 +- tests/test_galley_logical_optimizer.py | 2 +- tests/test_stats_interpreter.py | 9 ++++++--- tests/test_tensor_stats.py | 16 ++++++++++------ 23 files changed, 49 insertions(+), 37 deletions(-) create mode 100644 src/finchlite/adv_autoschedulers/__init__.py rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/__init__.py (68%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/blocked_stats.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/database_stats.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/dc_stats.py (99%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/dense_stat.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/dummy_stats.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/numeric_stats.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/stats_interpreter.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/tensor_def.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/tensor_stats.py (100%) rename src/finchlite/{autoschedule => adv_autoschedulers}/tensor_stats/uniform_stats.py (100%) diff --git a/src/finchlite/__init__.py b/src/finchlite/__init__.py index f050f8fb4..8e4e2dcc5 100644 --- a/src/finchlite/__init__.py +++ b/src/finchlite/__init__.py @@ -1,3 +1,13 @@ +from .adv_autoschedulers.tensor_stats import ( + BlockedStats, + DenseStats, + TensorDef, + UniformStats, +) +from .adv_autoschedulers.tensor_stats.dc_stats import ( + DC, + DCStats, +) from .algebra import ( FTyped, Tensor, @@ -30,14 +40,6 @@ uint32, uint64, ) -from .autoschedule.tensor_stats import ( - DC, - BlockedStats, - DCStats, - DenseStats, - TensorDef, - UniformStats, -) from .codegen import ( NumpyBuffer, NumpyBufferFType, diff --git a/src/finchlite/adv_autoschedulers/__init__.py b/src/finchlite/adv_autoschedulers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/finchlite/autoschedule/tensor_stats/__init__.py b/src/finchlite/adv_autoschedulers/tensor_stats/__init__.py similarity index 68% rename from src/finchlite/autoschedule/tensor_stats/__init__.py rename to src/finchlite/adv_autoschedulers/tensor_stats/__init__.py index 78d0a71d9..ecff1f943 100644 --- a/src/finchlite/autoschedule/tensor_stats/__init__.py +++ b/src/finchlite/adv_autoschedulers/tensor_stats/__init__.py @@ -1,27 +1,27 @@ from ...finch_logic import TensorStats from .blocked_stats import BlockedStats, BlockedStatsFactory -from .database_stats import DatabaseStats, DatabaseStatsFactory -from .dc_stats import DC, DCStats, DCStatsFactory from .dense_stat import DenseStats, DenseStatsFactory from .dummy_stats import DummyStats, DummyStatsFactory +from .numeric_stats import NumericStats from .stats_interpreter import StatsInterpreter from .tensor_def import TensorDef from .tensor_stats import BaseTensorStats from .uniform_stats import UniformStats, UniformStatsFactory +# dc_stats and database_stats are NOT imported here because they +# statically depend on compile/, which must be fully loaded first. +# Import DC, DCStats, DCStatsFactory, DatabaseStats, DatabaseStatsFactory +# directly from their submodules when needed. + __all__ = [ - "DC", "BaseTensorStats", "BlockedStats", "BlockedStatsFactory", - "DCStats", - "DCStatsFactory", - "DatabaseStats", - "DatabaseStatsFactory", "DenseStats", "DenseStatsFactory", "DummyStats", "DummyStatsFactory", + "NumericStats", "StatsInterpreter", "TensorDef", "TensorStats", diff --git a/src/finchlite/autoschedule/tensor_stats/blocked_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/blocked_stats.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/blocked_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/blocked_stats.py diff --git a/src/finchlite/autoschedule/tensor_stats/database_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/database_stats.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/database_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/database_stats.py diff --git a/src/finchlite/autoschedule/tensor_stats/dc_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/dc_stats.py similarity index 99% rename from src/finchlite/autoschedule/tensor_stats/dc_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/dc_stats.py index b95f3d742..a7216f173 100644 --- a/src/finchlite/autoschedule/tensor_stats/dc_stats.py +++ b/src/finchlite/adv_autoschedulers/tensor_stats/dc_stats.py @@ -13,6 +13,7 @@ from ... import finch_notation as ntn from ...algebra import Tensor, ffuncs, ftype, int64 from ...algebra.algebra import FinchOperator +from ...compile import BufferizedNDArray, make_extent from .numeric_stats import NumericStats from .tensor_def import TensorDef from .tensor_stats import BaseTensorStatsFactory @@ -148,8 +149,6 @@ def _structure_to_dcs(self, arr: Tensor, fields: Iterable[Field]) -> set[DC]: # For each field i, we compute DC({}, {i}) and DC({i}, {*fields}). # Additionally, we compute the nnz for the full tensor DC({}, {*fields}). def _array_to_dcs(self, arr: Any, fields: Iterable[Field]) -> set[DC]: - from ...compile import BufferizedNDArray, make_extent - int64_vector_ftype = BufferizedNDArray.from_numpy( np.zeros(1, dtype=np.int64) ).ftype diff --git a/src/finchlite/autoschedule/tensor_stats/dense_stat.py b/src/finchlite/adv_autoschedulers/tensor_stats/dense_stat.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/dense_stat.py rename to src/finchlite/adv_autoschedulers/tensor_stats/dense_stat.py diff --git a/src/finchlite/autoschedule/tensor_stats/dummy_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/dummy_stats.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/dummy_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/dummy_stats.py diff --git a/src/finchlite/autoschedule/tensor_stats/numeric_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/numeric_stats.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/numeric_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/numeric_stats.py diff --git a/src/finchlite/autoschedule/tensor_stats/stats_interpreter.py b/src/finchlite/adv_autoschedulers/tensor_stats/stats_interpreter.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/stats_interpreter.py rename to src/finchlite/adv_autoschedulers/tensor_stats/stats_interpreter.py diff --git a/src/finchlite/autoschedule/tensor_stats/tensor_def.py b/src/finchlite/adv_autoschedulers/tensor_stats/tensor_def.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/tensor_def.py rename to src/finchlite/adv_autoschedulers/tensor_stats/tensor_def.py diff --git a/src/finchlite/autoschedule/tensor_stats/tensor_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/tensor_stats.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/tensor_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/tensor_stats.py diff --git a/src/finchlite/autoschedule/tensor_stats/uniform_stats.py b/src/finchlite/adv_autoschedulers/tensor_stats/uniform_stats.py similarity index 100% rename from src/finchlite/autoschedule/tensor_stats/uniform_stats.py rename to src/finchlite/adv_autoschedulers/tensor_stats/uniform_stats.py diff --git a/src/finchlite/autoschedule/cache.py b/src/finchlite/autoschedule/cache.py index d8d775393..468965fa1 100644 --- a/src/finchlite/autoschedule/cache.py +++ b/src/finchlite/autoschedule/cache.py @@ -3,8 +3,8 @@ import numpy as np from numpy.linalg import vector_norm +from finchlite.adv_autoschedulers.tensor_stats.numeric_stats import NumericStats from finchlite.algebra.tensor import TensorFType -from finchlite.autoschedule.tensor_stats.numeric_stats import NumericStats from finchlite.finch_logic import ( Alias, LogicLoader, diff --git a/src/finchlite/autoschedule/einsum.py b/src/finchlite/autoschedule/einsum.py index 22fe9e3c9..5a7f514f2 100644 --- a/src/finchlite/autoschedule/einsum.py +++ b/src/finchlite/autoschedule/einsum.py @@ -2,10 +2,9 @@ import finchlite.finch_logic as lgc from finchlite.algebra import ffuncs, ftype, init_value from finchlite.algebra.tensor import TensorFType -from finchlite.autoschedule.tensor_stats import TensorStats from finchlite.finch_assembly.stages import AssemblyLibrary from finchlite.finch_einsum import EinsumLoader, MockEinsumLoader -from finchlite.finch_logic import LogicStatement, StatsFactory +from finchlite.finch_logic import LogicStatement, StatsFactory, TensorStats from finchlite.finch_logic.stages import LogicLoader from .stages import LogicEinsumLowerer diff --git a/src/finchlite/autoschedule/executor.py b/src/finchlite/autoschedule/executor.py index 3dd17615c..cf434c3bc 100644 --- a/src/finchlite/autoschedule/executor.py +++ b/src/finchlite/autoschedule/executor.py @@ -5,7 +5,6 @@ from finchlite.finch_logic.nodes import TableValue from .. import finch_logic as lgc -from ..autoschedule.tensor_stats import DenseStatsFactory from ..finch_logic import ( LogicEvaluator, LogicLoader, @@ -59,6 +58,8 @@ def __init__( if ctx is None: ctx = DefaultLogicFormatter() if stats_factory is None: + from ..adv_autoschedulers.tensor_stats import DenseStatsFactory + stats_factory = DenseStatsFactory() self.ctx: LogicLoader = ctx self.stats_factory = stats_factory diff --git a/src/finchlite/autoschedule/galley/logical_optimizer/annotated_query.py b/src/finchlite/autoschedule/galley/logical_optimizer/annotated_query.py index 9de765ec1..1e08b08b3 100644 --- a/src/finchlite/autoschedule/galley/logical_optimizer/annotated_query.py +++ b/src/finchlite/autoschedule/galley/logical_optimizer/annotated_query.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import Any, cast +from ....adv_autoschedulers.tensor_stats.numeric_stats import NumericStats from ....algebra import ( cansplitpush, ffuncs, @@ -36,7 +37,6 @@ intree, isdescendant, ) -from ...tensor_stats.numeric_stats import NumericStats from .logic_to_stats import insert_statistics diff --git a/src/finchlite/interface/lazy.py b/src/finchlite/interface/lazy.py index 1c620d89c..9b8e1e4da 100644 --- a/src/finchlite/interface/lazy.py +++ b/src/finchlite/interface/lazy.py @@ -27,7 +27,6 @@ return_type, ) from ..algebra.ftypes import FDType -from ..autoschedule.tensor_stats import StatsInterpreter from ..finch_logic import ( Aggregate, Alias, @@ -2129,6 +2128,8 @@ def get_lazy_tensor_stats( lazy_tensor: LazyTensor, stats_factory: StatsFactory, ) -> TensorStats: + from ..adv_autoschedulers.tensor_stats import StatsInterpreter + trace = lazy_tensor.ctx.trace() interpreter = StatsInterpreter(stats_factory=stats_factory) bindings: OrderedDict[Alias, TensorStats] = OrderedDict() diff --git a/tests/test_cache_loader.py b/tests/test_cache_loader.py index fff61cc27..a85e5b5d2 100644 --- a/tests/test_cache_loader.py +++ b/tests/test_cache_loader.py @@ -3,14 +3,17 @@ import numpy as np import finchlite as fl +from finchlite.adv_autoschedulers.tensor_stats import ( + DenseStatsFactory, + UniformStatsFactory, +) +from finchlite.adv_autoschedulers.tensor_stats.blocked_stats import BlockedStatsFactory +from finchlite.adv_autoschedulers.tensor_stats.dc_stats import DCStatsFactory from finchlite.algebra import ffuncs from finchlite.autoschedule.cache import ( LogicCacheLRU_Embeddings_Norms, ) from finchlite.autoschedule.executor import LogicExecutor -from finchlite.autoschedule.tensor_stats import DenseStatsFactory, UniformStatsFactory -from finchlite.autoschedule.tensor_stats.blocked_stats import BlockedStatsFactory -from finchlite.autoschedule.tensor_stats.dc_stats import DCStatsFactory from finchlite.finch_logic import ( Alias, Field, diff --git a/tests/test_galley_branch_and_bound.py b/tests/test_galley_branch_and_bound.py index f0070f950..b42f41e19 100644 --- a/tests/test_galley_branch_and_bound.py +++ b/tests/test_galley_branch_and_bound.py @@ -7,6 +7,7 @@ import numpy as np import finchlite as fl +from finchlite.adv_autoschedulers.tensor_stats import DenseStatsFactory from finchlite.algebra import ffuncs from finchlite.autoschedule.galley.logical_optimizer import AnnotatedQuery from finchlite.autoschedule.galley.logical_optimizer.branch_and_bound import ( @@ -14,7 +15,6 @@ branch_and_bound_dfs, pruned_query_to_plan, ) -from finchlite.autoschedule.tensor_stats import DenseStatsFactory from finchlite.finch_logic import ( Aggregate, Alias, diff --git a/tests/test_galley_logical_optimizer.py b/tests/test_galley_logical_optimizer.py index 3055c1509..ca5ac1d8b 100644 --- a/tests/test_galley_logical_optimizer.py +++ b/tests/test_galley_logical_optimizer.py @@ -6,12 +6,12 @@ import finchlite as fl from finchlite import ffuncs +from finchlite.adv_autoschedulers.tensor_stats import DenseStatsFactory from finchlite.autoschedule.galley.logical_optimizer import ( AnnotatedQuery, greedy_query, insert_statistics, ) -from finchlite.autoschedule.tensor_stats import DenseStatsFactory from finchlite.finch_logic import ( Aggregate, Alias, diff --git a/tests/test_stats_interpreter.py b/tests/test_stats_interpreter.py index c710f7a22..ecd4f0550 100644 --- a/tests/test_stats_interpreter.py +++ b/tests/test_stats_interpreter.py @@ -3,12 +3,15 @@ import numpy as np import finchlite as fl -from finchlite.algebra import ffuncs -from finchlite.autoschedule.tensor_stats import DatabaseStatsFactory, DCStatsFactory -from finchlite.autoschedule.tensor_stats.stats_interpreter import ( +from finchlite.adv_autoschedulers.tensor_stats.database_stats import ( + DatabaseStatsFactory, +) +from finchlite.adv_autoschedulers.tensor_stats.dc_stats import DCStatsFactory +from finchlite.adv_autoschedulers.tensor_stats.stats_interpreter import ( StatsInterpreter, calculate_estimated_error, ) +from finchlite.algebra import ffuncs from finchlite.finch_logic import ( Aggregate, Alias, diff --git a/tests/test_tensor_stats.py b/tests/test_tensor_stats.py index cf8d2e3e4..a2b23e42d 100644 --- a/tests/test_tensor_stats.py +++ b/tests/test_tensor_stats.py @@ -7,14 +7,9 @@ import finchlite as fl from finchlite import ffuncs -from finchlite.autoschedule.galley.logical_optimizer import insert_statistics -from finchlite.autoschedule.tensor_stats import ( - DC, +from finchlite.adv_autoschedulers.tensor_stats import ( BlockedStats, BlockedStatsFactory, - DatabaseStatsFactory, - DCStats, - DCStatsFactory, DenseStats, DenseStatsFactory, DummyStatsFactory, @@ -22,6 +17,15 @@ UniformStats, UniformStatsFactory, ) +from finchlite.adv_autoschedulers.tensor_stats.database_stats import ( + DatabaseStatsFactory, +) +from finchlite.adv_autoschedulers.tensor_stats.dc_stats import ( + DC, + DCStats, + DCStatsFactory, +) +from finchlite.autoschedule.galley.logical_optimizer import insert_statistics from finchlite.finch_logic import ( Aggregate, Field,