Dense Level

src/finchlite/__init__.py

@@ -135,6 +135,9 @@
     DenseLevelFType,
     ElementLevelFType,
     FiberTensorFType,
+    dense,
+    element,
+    fiber_tensor,
 )
 
 __all__ = [
@@ -188,17 +191,20 @@
     "cos",
     "cosh",
     "defer",
+    "dense",
     "dimension",
     "divide",
     "einop",
     "einsum",
+    "element",
     "element_type",
     "elementwise",
     "equal",
     "exp",
     "expand_dims",
     "expm1",
     "extent",
+    "fiber_tensor",
     "fill_value",
     "fisinstance",
     "flatten",

src/finchlite/algebra/tensor.py

@@ -34,6 +34,11 @@ def shape_type(self) -> tuple[type, ...]:
         e.g. dtypes, formats, or types, and so that we can easily index it."""
         ...
 
+    @abstractmethod
+    def __init__(self, *args):
+        """TensorFType instance initializer."""
+        ...
+
 
 class Tensor(FTyped, ABC):
     """

src/finchlite/autoschedule/compiler.py

@@ -14,8 +14,7 @@
     return_type,
 )
 from ..codegen import NumpyBufferFType
-from ..compile import BufferizedNDArrayFType, ExtentFType, dimension
-from ..finch_assembly import TupleFType
+from ..compile import ExtentFType, dimension
 from ..finch_logic import (
     Aggregate,
     Alias,
@@ -204,11 +203,7 @@ def __call__(
                 return ntn.Assign(
                     ntn.Variable(
                         name,
-                        BufferizedNDArrayFType(
-                            NumpyBufferFType(val.dtype),
-                            val.ndim,
-                            TupleFType.from_tuple(val.shape_type),
-                        ),
+                        val.from_kwargs(val.to_kwargs()),
                     ),
                     compile_logic_constant(tns),
                 )
@@ -484,13 +479,19 @@ def find_suitable_rep(root, table_vars) -> TensorFType:
                 )
             )
 
-            return BufferizedNDArrayFType(
-                buf_t=NumpyBufferFType(dtype),
+            # TODO: properly infer result rep from args
+            result_rep, fields = args_suitable_reps_fields[0]
+            levels_to_add = [
+                idx for idx, f in enumerate(result_fields) if f not in fields
+            ]
+            result_rep = result_rep.add_levels(levels_to_add)
+            kwargs = result_rep.to_kwargs()
+            kwargs.update(
+                element_type=NumpyBufferFType(dtype),
                 ndim=np.intp(len(result_fields)),
-                strides_t=TupleFType.from_tuple(
-                    tuple(field_type_map[f] for f in result_fields)
-                ),
+                dimension_type=tuple(field_type_map[f] for f in result_fields),
             )
+            return result_rep.from_kwargs(**kwargs)
         case Aggregate(Literal(op), init, arg, idxs):
             init_suitable_rep = find_suitable_rep(init, table_vars)
             arg_suitable_rep = find_suitable_rep(arg, table_vars)
@@ -499,16 +500,24 @@ def find_suitable_rep(root, table_vars) -> TensorFType:
                     op, init_suitable_rep.element_type, arg_suitable_rep.element_type
                 )
             )
-            strides_t = tuple(
-                st
-                for f, st in zip(arg.fields, arg_suitable_rep.shape_type, strict=True)
-                if f not in idxs
-            )
-            return BufferizedNDArrayFType(
-                buf_t=buf_t,
+            # TODO: properly infer result rep from args
+            levels_to_remove = []
+            strides_t = []
+            for idx, (f, st) in enumerate(
+                zip(arg.fields, arg_suitable_rep.shape_type, strict=True)
+            ):
+                if f not in idxs:
+                    strides_t.append(st)
+                else:
+                    levels_to_remove.append(idx)
+            arg_suitable_rep = arg_suitable_rep.remove_levels(levels_to_remove)
+            kwargs = arg_suitable_rep.to_kwargs()
+            kwargs.update(
+                buffer_type=buf_t,
                 ndim=np.intp(len(strides_t)),
-                strides_t=TupleFType.from_tuple(strides_t),
+                dimension_type=tuple(strides_t),
             )
+            return arg_suitable_rep.from_kwargs(**kwargs)
         case LogicTree() as tree:
             for child in tree.children:
                 suitable_rep = find_suitable_rep(child, table_vars)
@@ -548,4 +557,4 @@ def __call__(self, prgm: LogicNode) -> tuple[ntn.NotationNode, dict[Alias, Table
         lowered_prgm = self.ll(
             prgm, table_vars, slot_vars, dim_size_vars, field_relabels
         )
-        return merge_blocks(lowered_prgm), tables
+        return merge_blocks(lowered_prgm), table_vars, tables

src/finchlite/autoschedule/optimize.py

@@ -321,7 +321,7 @@ def propagate_into_reformats(root: LogicNode) -> LogicNode:
     class Entry:
         node: Query
         node_pos: int
-        matched: Query[LogicNode, Reformat] | None = None
+        matched: Query | None = None
         matched_pos: int | None = None
 
     def rule_0(ex: LogicNode) -> LogicNode | None:
@@ -346,6 +346,7 @@ def rule_0(ex: LogicNode) -> LogicNode | None:
                         if q.node.lhs not in PostOrderDFS(
                             Plan(tuple(new_bodies[q.node_pos + 1 :]))
                         ) and isinstance(q.node.rhs, MapJoin | Aggregate | Reorder):
+                            assert isinstance(q.matched.rhs, Reformat)
                             new_bodies[q.node_pos] = Query(
                                 q.matched.lhs, Reformat(q.matched.rhs.tns, q.node.rhs)
                             )

src/finchlite/codegen/numba_backend.py

@@ -716,11 +716,11 @@ def struct_numba_setattr(fmt: AssemblyStructFType, ctx, obj, attr, val):
 
 
 def struct_construct_from_numba(fmt: AssemblyStructFType, numba_struct):
-    args = [
-        construct_from_numba(field_type, getattr(numba_struct, name))
+    kwargs = {
+        name: construct_from_numba(field_type, getattr(numba_struct, name))
         for (name, field_type) in fmt.struct_fields
-    ]
-    return fmt(*args)
+    }
+    return fmt(**kwargs)
 
 
 register_property(

src/finchlite/codegen/numpy_buffer.py

@@ -76,6 +76,10 @@ def __str__(self):
         arr_str = str(self.arr).replace("\n", "")
         return f"np_buf({arr_str})"
 
+    def __repr__(self):
+        arr_repr = repr(self.arr).replace("\n", "")
+        return f"NumpyBuffer({arr_repr})"
+
 
 class NumpyBufferFType(CBufferFType, NumbaBufferFType, CStackFType):
     """

src/finchlite/compile/bufferized_ndarray.py

@@ -16,24 +16,24 @@
 class BufferizedNDArray(Tensor):
     def __init__(
         self,
-        arr: np.ndarray | NumpyBuffer,
+        val: np.ndarray | NumpyBuffer,
         shape: tuple[np.integer, ...] | None = None,
         strides: tuple[np.integer, ...] | None = None,
     ):
         self._shape: tuple[np.integer, ...]
         self.strides: tuple[np.integer, ...]
-        if shape is None and strides is None and isinstance(arr, np.ndarray):
-            itemsize = arr.dtype.itemsize
-            for stride in arr.strides:
+        if shape is None and strides is None and isinstance(val, np.ndarray):
+            itemsize = val.dtype.itemsize
+            for stride in val.strides:
                 if stride % itemsize != 0:
                     raise ValueError("Array must be aligned to multiple of itemsize")
-            self.strides = tuple(np.intp(stride // itemsize) for stride in arr.strides)
-            self._shape = tuple(np.intp(s) for s in arr.shape)
-            self.buf = NumpyBuffer(arr.reshape(-1, copy=False))
-        elif shape is not None and strides is not None and isinstance(arr, NumpyBuffer):
+            self.strides = tuple(np.intp(stride // itemsize) for stride in val.strides)
+            self._shape = tuple(np.intp(s) for s in val.shape)
+            self.val = NumpyBuffer(val.reshape(-1, copy=False))
+        elif shape is not None and strides is not None and isinstance(val, NumpyBuffer):
             self.strides = strides
             self._shape = shape
-            self.buf = arr
+            self.val = val
         else:
             raise Exception("Invalid constructor arguments")
 
@@ -42,14 +42,18 @@ def to_numpy(self):
         Convert the bufferized NDArray to a NumPy array.
         This is used to get the underlying NumPy array from the bufferized NDArray.
         """
-        return self.buf.arr.reshape(self._shape, copy=False)
+        return self.val.arr.reshape(self._shape, copy=False)
 
     @property
     def ftype(self):
         """
         Returns the ftype of the buffer, which is a BufferizedNDArrayFType.
         """
-        return BufferizedNDArrayFType(ftype(self.buf), self.ndim, ftype(self.strides))
+        return BufferizedNDArrayFType(
+            buffer_type=ftype(self.val),
+            ndim=self.ndim,
+            dimension_type=ftype(self.strides),
+        )
 
     @property
     def shape(self):
@@ -74,8 +78,8 @@ def declare(self, init, op, shape):
                 raise ValueError(
                     f"Invalid dimension end value {dim.end} for ndarray declaration."
                 )
-        for i in range(self.buf.length()):
-            self.buf.store(i, init)
+        for i in range(self.val.length()):
+            self.val.store(i, init)
         return self
 
     def freeze(self, op):
@@ -103,7 +107,7 @@ def __setitem__(self, index, value):
         """
         if isinstance(index, tuple):
             index = np.ravel_multi_index(index, self._shape)
-        self.buf.store(index, value)
+        self.val.store(index, value)
 
     def __str__(self):
         return f"BufferizedNDArray(shape={self.shape})"
@@ -141,24 +145,30 @@ def str_format(types):
     @property
     def struct_fields(self):
         return [
-            ("buf", self.buf_t),
+            ("val", self.buf_t),
             ("shape", self.shape_t),
             ("strides", self.strides_t),
         ]
 
-    def __init__(self, buf_t: NumpyBufferFType, ndim: np.intp, strides_t: TupleFType):
-        self.buf_t = buf_t
+    def __init__(
+        self,
+        *,
+        buffer_type: NumpyBufferFType,
+        ndim: np.intp,
+        dimension_type: TupleFType,
+    ):
+        self.buf_t = buffer_type
         self._ndim = ndim
-        self.shape_t = strides_t  # assuming shape is the same type as strides
-        self.strides_t = strides_t
+        self.shape_t = dimension_type  # assuming shape is the same type as strides
+        self.strides_t = dimension_type
 
     def __eq__(self, other):
         if not isinstance(other, BufferizedNDArrayFType):
             return False
-        return self.buf_t == other.buf_t and self._ndim == other._ndim
+        return self.buf_t == other.buf_t and self.ndim == other.ndim
 
     def __hash__(self):
-        return hash((self.buf_t, self._ndim))
+        return hash((self.buf_t, self.ndim))
 
     def __str__(self):
         return str(self.struct_name)
@@ -170,6 +180,35 @@ def __repr__(self):
     def ndim(self) -> np.intp:
         return self._ndim
 
+    @ndim.setter
+    def ndim(self, val):
+        self._ndim = val
+
+    def from_kwargs(self, **kwargs) -> "BufferizedNDArrayFType":
+        b_t = kwargs.get("buffer_type", self.buf_t)
+        ndim = kwargs.get("ndim", self.ndim)
+        if "shape_type" in kwargs:
+            s_t = kwargs["shape_type"]
+            d_t = s_t if isinstance(s_t, TupleFType) else TupleFType.from_tuple(s_t)
+        else:
+            d_t = self.shape_t
+        return BufferizedNDArrayFType(buffer_type=b_t, ndim=ndim, dimension_type=d_t)
+
+    def to_kwargs(self):
+        return {
+            "buffer_type": self.buf_t,
+            "ndim": self.ndim,
+            "shape_type": self.shape_t,
+        }
+
+    # TODO: temporary approach for suitable rep and traits
+    def add_levels(self, idxs: list[int]):
+        return self
+
+    # TODO: temporary approach for suitable rep and traits
+    def remove_levels(self, idxs: list[int]):
+        return self
+
     @property
     def fill_value(self) -> Any:
         return np.zeros((), dtype=self.buf_t.element_type)[()]
@@ -180,7 +219,7 @@ def element_type(self):
 
     @property
     def shape_type(self) -> tuple:
-        return tuple(np.intp for _ in range(self._ndim))
+        return tuple(np.intp for _ in range(self.ndim))
 
     def lower_declare(self, ctx, tns, init, op, shape):
         i_var = asm.Variable("i", self.buf_t.length_type)
@@ -220,14 +259,14 @@ def asm_unpack(self, ctx, var_n, val):
         Unpack the into asm context.
         """
         stride = []
-        for i in range(self._ndim):
+        for i in range(self.ndim):
             stride_i = asm.Variable(f"{var_n}_stride_{i}", self.buf_t.length_type)
             stride.append(stride_i)
             stride_e = asm.GetAttr(val, asm.Literal("strides"))
             stride_i_e = asm.GetAttr(stride_e, asm.Literal(f"element_{i}"))
             ctx.exec(asm.Assign(stride_i, stride_i_e))
         buf = asm.Variable(f"{var_n}_buf", self.buf_t)
-        buf_e = asm.GetAttr(val, asm.Literal("buf"))
+        buf_e = asm.GetAttr(val, asm.Literal("val"))
         ctx.exec(asm.Assign(buf, buf_e))
         buf_s = asm.Slot(f"{var_n}_buf_slot", self.buf_t)
         ctx.exec(asm.Unpack(buf_s, buf))
@@ -243,11 +282,11 @@ def asm_repack(self, ctx, lhs, obj):
 
     def __call__(
         self,
-        buf: NumpyBuffer,
-        shape: tuple[np.integer, ...],
-        strides: tuple[np.integer, ...],
+        val: NumpyBuffer,
+        shape: tuple[np.integer, ...] | None = None,
+        strides: tuple[np.integer, ...] | None = None,
     ) -> BufferizedNDArray:
-        return BufferizedNDArray(buf, shape, strides)
+        return BufferizedNDArray(val, shape, strides)
 
 
 class BufferizedNDArrayAccessor(Tensor):
@@ -288,7 +327,7 @@ def unwrap(self):
         This is used to get the original tensor from a tensor view.
         """
         assert self.ndim == 0, "Cannot unwrap a tensor view with non-zero dimension."
-        return self.tns.buf.load(self.pos)
+        return self.tns.val.load(self.pos)
 
     def increment(self, val):
         """
@@ -298,7 +337,7 @@ def increment(self, val):
         if self.op is None:
             raise ValueError("No operation defined for increment.")
         assert self.ndim == 0, "Cannot unwrap a tensor view with non-zero dimension."
-        self.tns.buf.store(self.pos, self.op(self.tns.buf.load(self.pos), val))
+        self.tns.val.store(self.pos, self.op(self.tns.val.load(self.pos), val))
         return self
 
 
@@ -377,7 +416,7 @@ def asm_repack(self, ctx, lhs, obj):
         """
         Repack the buffer from C context.
         """
-        (self.tns.asm_repack(ctx, lhs.tns, obj.tns),)
+        self.tns.asm_repack(ctx, lhs.tns, obj.tns)
         ctx.exec(
             asm.Block(
                 asm.SetAttr(lhs, "tns", obj.tns),