Skip to content

Commit

Permalink
Fix serialization issues
Browse files Browse the repository at this point in the history
  • Loading branch information
@JanKleine
JanKleine committed Oct 26, 2023
1 parent 33aa5b3 commit 7586113
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 34 deletions.
98 changes: 72 additions & 26 deletions dace/data.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -624,7 +624,38 @@ def fields(self, lvl: int, dummy_symbol: symbolic.SymExpr) -> Dict[str, Data]:
pass


class Dense(TensorIndex):
def to_json(self):
attrs = serialize.all_properties_to_json(self)

retdict = {"type": type(self).__name__, "attributes": attrs}

return retdict


@classmethod
def from_json(cls, json_obj, context=None):

# Selecting proper subclass
if json_obj['type'] == "TensorIndexDense":
self = TensorIndexDense.__new__(TensorIndexDense)
elif json_obj['type'] == "TensorIndexCompressed":
self = TensorIndexCompressed.__new__(TensorIndexCompressed)
elif json_obj['type'] == "TensorIndexSingleton":
self = TensorIndexSingleton.__new__(TensorIndexSingleton)
elif json_obj['type'] == "TensorIndexRange":
self = TensorIndexRange.__new__(TensorIndexRange)
elif json_obj['type'] == "TensorIndexOffset":
self = TensorIndexOffset.__new__(TensorIndexOffset)
else:
raise TypeError(f"Invalid data type, got: {json_obj['type']}")

serialize.set_properties_from_json(self, json_obj['attributes'], context=context)

return self


@make_properties
class TensorIndexDense(TensorIndex):
"""
Dense tensor index.
Expand All @@ -633,8 +664,8 @@ class Dense(TensorIndex):
index structure beyond the corresponding dimension size.
"""

_ordered: bool
_unique: bool
_ordered = Property(dtype=bool, default=False)
_unique = Property(dtype=bool)

@property
def iteration_type(self) -> TensorIterationTypes:
Expand Down Expand Up @@ -690,7 +721,8 @@ def __repr__(self) -> str:
return s


class Compressed(TensorIndex):
@make_properties
class TensorIndexCompressed(TensorIndex):
"""
Tensor level that stores coordinates in segmented array.
Expand All @@ -699,9 +731,9 @@ class Compressed(TensorIndex):
array that holds the child coordinates corresponding the parent.
"""

_full: bool
_ordered: bool
_unique: bool
_full = Property(dtype=bool, default=False)
_ordered = Property(dtype=bool, default=False)
_unique = Property(dtype=bool, default=False)

@property
def iteration_type(self) -> TensorIterationTypes:
Expand Down Expand Up @@ -766,7 +798,8 @@ def __repr__(self) -> str:
return s


class Singleton(TensorIndex):
@make_properties
class TensorIndexSingleton(TensorIndex):
"""
Tensor index that encodes a single coordinate per parent coordinate.
Expand All @@ -775,9 +808,9 @@ class Singleton(TensorIndex):
coordinate but the first is encoded in this manner.
"""

_full: bool
_ordered: bool
_unique: bool
_full = Property(dtype=bool, default=False)
_ordered = Property(dtype=bool, default=False)
_unique = Property(dtype=bool, default=False)

@property
def iteration_type(self) -> TensorIterationTypes:
Expand Down Expand Up @@ -841,7 +874,8 @@ def __repr__(self) -> str:
return s


class Range(TensorIndex):
@make_properties
class TensorIndexRange(TensorIndex):
"""
Tensor index that encodes a interval of coordinates for every parent.
Expand All @@ -851,8 +885,8 @@ class Range(TensorIndex):
range of coordinates between max(0, -offset[i]) and min(N, M - offset[i]).
"""

_ordered: bool
_unique: bool
_ordered = Property(dtype=bool, default=False)
_unique = Property(dtype=bool, default=False)

@property
def iteration_type(self) -> TensorIterationTypes:
Expand Down Expand Up @@ -910,16 +944,17 @@ def __repr__(self) -> str:
return s


class Offset(TensorIndex):
@make_properties
class TensorIndexOffset(TensorIndex):
"""
Tensor index that encodes the next coordinates as offset from parent.
Given a parent coordinate i and an offset index k, the level encodes the
coordinate j = i + offset[k].
"""

_ordered: bool
_unique: bool
_ordered = Property(dtype=bool, default=False)
_unique = Property(dtype=bool, default=False)

@property
def iteration_type(self) -> TensorIterationTypes:
Expand Down Expand Up @@ -987,14 +1022,15 @@ class Tensor(Structure):

value_dtype = TypeClassProperty(default=dtypes.int32, choices=dtypes.Typeclasses)
tensor_shape = ShapeProperty(default=[])
indices = ListProperty(element_type=Union[Tuple[TensorIndex, int], Tuple[TensorIndex, symbolic.SymExpr]])
value_count = SymbolicProperty
indices = ListProperty(element_type=TensorIndex)
index_ordering = ListProperty(element_type=symbolic.SymExpr)
value_count = SymbolicProperty(default=0)

def __init__(
self,
value_dtype: dtypes.Typeclasses,
tensor_shape,
indices: List[Union[Tuple[TensorIndex, int], Tuple[TensorIndex, symbolic.SymExpr]]],
indices: List[Tuple[TensorIndex, Union[int, symbolic.SymExpr]]],
value_count: symbolic.SymExpr,
name: str,
transient: bool = False,
Expand Down Expand Up @@ -1100,11 +1136,13 @@ def __init__(

self.value_dtype = value_dtype
self.tensor_shape = tensor_shape
self.indices = indices
self.value_count = value_count

indices, index_ordering = zip(*indices)
self.indices, self.index_ordering = list(indices), list(index_ordering)

num_dims = len(tensor_shape)
dimension_order = [idx[1] for idx in indices if isinstance(idx[1], int)]
dimension_order = [idx for idx in self.index_ordering if isinstance(idx, int)]

# all tensor dimensions must occure exactly once in indices
if not sorted(dimension_order) == list(range(num_dims)):
Expand All @@ -1122,18 +1160,26 @@ def __init__(
values=dtypes.float32[value_count],
)

indices_pure = [idx[0] for idx in indices];
for (lvl, index) in enumerate(indices_pure):
for (lvl, index) in enumerate(indices):
fields.update(index.fields(lvl, value_count))

abbreviation = ''.join(str(idx)[0] for idx in indices_pure)

super(Tensor, self).__init__(fields, name, transient, storage, location,
lifetime, debuginfo)

def __repr__(self):
return f"{self.name} (dtype: {self.value_dtype}, shape: {list(self.tensor_shape)}, indices: {self.indices})"

@staticmethod
def from_json(json_obj, context=None):
if json_obj['type'] != 'Tensor':
raise TypeError("Invalid data type")

# Create dummy object
tensor = Tensor.__new__(Tensor)
serialize.set_properties_from_json(tensor, json_obj, context=context)

return tensor


@make_properties
class StructureView(Structure):
Expand Down
18 changes: 10 additions & 8 deletions tests/sdfg/data/tensor_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ def test_read_csr_tensor():
csr_obj = dace.data.Tensor(
dace.float32,
(M, N),
[(dace.data.Dense(), 0), (dace.data.Compressed(), 1)],
[(dace.data.TensorIndexDense(), 0), (dace.data.TensorIndexCompressed(), 1)],
nnz,
"CSR_Tensor")

Expand Down Expand Up @@ -63,6 +63,8 @@ def test_read_csr_tensor():
func(A=inpA, B=B, M=A.shape[0], N=A.shape[1], nnz=A.nnz)
ref = A.toarray()

sdfg.save("./tensor.json")

assert np.allclose(B, ref)


Expand All @@ -73,7 +75,7 @@ def test_csr_fields():
csr = dace.data.Tensor(
dace.float32,
(M, N),
[(dace.data.Dense(), 0), (dace.data.Compressed(), 1)],
[(dace.data.TensorIndexDense(), 0), (dace.data.TensorIndexCompressed(), 1)],
nnz,
"CSR_Matrix",
)
Expand All @@ -90,9 +92,9 @@ def test_dia_fields():
dace.float32,
(M, N),
[
(dace.data.Dense(), num_diags),
(dace.data.Range(), 0),
(dace.data.Offset(), 1),
(dace.data.TensorIndexDense(), num_diags),
(dace.data.TensorIndexRange(), 0),
(dace.data.TensorIndexOffset(), 1),
],
nnz,
"DIA_Matrix",
Expand All @@ -110,9 +112,9 @@ def test_coo_fields():
dace.float32,
(I, J, K),
[
(dace.data.Compressed(unique=False), 0),
(dace.data.Singleton(unique=False), 1),
(dace.data.Singleton(), 2),
(dace.data.TensorIndexCompressed(unique=False), 0),
(dace.data.TensorIndexSingleton(unique=False), 1),
(dace.data.TensorIndexSingleton(), 2),
],
nnz,
"COO_3D_Tensor",
Expand Down

0 comments on commit 7586113

Please sign in to comment.