molcraft 0.1.0a1__py3-none-any.whl

molcraft/records.py

@@ -0,0 +1,169 @@
+import os
+import math
+import glob
+import time
+import typing 
+import tensorflow as tf
+import numpy as np
+import pandas as pd
+import multiprocessing as mp
+
+from molcraft import tensors
+from molcraft import featurizers
+
+
+def write(
+    inputs: list[str | tuple], 
+    featurizer: featurizers.Featurizer,
+    path: str, 
+    overwrite: bool = True, 
+    num_files: typing.Optional[int] = None, 
+    num_processes: typing.Optional[int] = None,
+    multiprocessing: bool = False,
+    device: str = '/cpu:0'
+) -> None:
+    
+    if os.path.isdir(path) and not overwrite:
+        return 
+    
+    os.makedirs(path, exist_ok=True)
+
+    with tf.device(device):
+
+        if isinstance(inputs, (pd.DataFrame, pd.Series)):
+            inputs = inputs.values
+        
+        if not isinstance(inputs, list):
+            inputs = list(inputs)
+
+        example = _featurize_input(inputs[0], featurizer)
+        if not isinstance(example, tensors.GraphTensor):
+            example = example[0]
+
+        save_spec(os.path.join(path, 'spec.pb'), example.spec)
+
+        if num_processes is None:
+            num_processes = mp.cpu_count()
+
+        if num_files is None:
+            num_files = min(len(inputs), num_processes)
+            
+        chunk_size = math.ceil(len(inputs) / num_files)
+        num_files = math.ceil(len(inputs) / chunk_size)
+        
+        paths = [
+            os.path.join(path, f'tfrecord-{i:04d}.tfrecord')
+            for i in range(num_files)
+        ]
+        
+        input_chunks = [
+            inputs[i * chunk_size: (i + 1) * chunk_size]
+            for i in range(num_files)
+        ]
+        
+        if not multiprocessing:
+            for path, input_chunk in zip(paths, input_chunks):
+                _write_tfrecord(input_chunk, path, featurizer)
+            return
+        
+        processes = []
+        
+        for path, input_chunk in zip(paths, input_chunks):
+        
+            while len(processes) >= num_processes:
+                for process in processes:
+                    if not process.is_alive():
+                        processes.remove(process)
+                else:
+                    time.sleep(0.1)
+                    continue
+                    
+            process = mp.Process(
+                target=_write_tfrecord,
+                args=(input_chunk, path, featurizer)
+            )
+            processes.append(process)
+            process.start()
+
+        for process in processes:
+            process.join()         
+    
+def load(
+    path: str, 
+    shuffle_files: bool = False
+) -> tf.data.Dataset:
+    spec = load_spec(os.path.join(path, 'spec.pb'))
+    filenames = sorted(glob.glob(os.path.join(path, '*.tfrecord')))
+    num_files = len(filenames)
+    ds = tf.data.Dataset.from_tensor_slices(filenames)
+    if shuffle_files:
+        ds = ds.shuffle(num_files)
+    ds = ds.interleave(
+        tf.data.TFRecordDataset, num_parallel_calls=1)
+    ds = ds.map(
+        lambda x: _parse_example(x, spec),
+        num_parallel_calls=tf.data.AUTOTUNE)
+    if not tensors.is_scalar(spec):
+        ds = ds.unbatch()
+    return ds
+
+def _write_tfrecord(
+    inputs, 
+    path: str,
+    featurizer: featurizers.Featurizer, 
+) -> None:
+    
+    def write_example(tensor):
+        flat_values = tf.nest.flatten(tensor, expand_composites=True)
+        flat_values = [tf.io.serialize_tensor(value).numpy() for value in flat_values]
+        feature = tf.train.Feature(bytes_list=tf.train.BytesList(value=flat_values))
+        serialized_feature = _serialize_example({'feature': feature})
+        writer.write(serialized_feature)
+
+    with tf.io.TFRecordWriter(path) as writer:
+        for x in inputs:
+            tensor = _featurize_input(x, featurizer)
+            if isinstance(tensor, tensors.GraphTensor):
+                write_example(tensor)
+            else:
+                for t in tensor:
+                    write_example(t)
+
+def _featurize_input(x, featurizer):
+    if isinstance(x, (list, np.ndarray)):
+        x = tuple(x)
+    return featurizer(x)
+
+def _serialize_example(
+    feature: dict[str, tf.train.Feature]
+) -> bytes:
+    example_proto = tf.train.Example(
+        features=tf.train.Features(feature=feature))
+    return example_proto.SerializeToString()
+
+def _parse_example(
+    x: tf.Tensor, 
+    spec: tensors.GraphTensor.Spec
+) -> tf.Tensor:
+    out = tf.io.parse_single_example(
+        x, features={'feature': tf.io.RaggedFeature(tf.string)})['feature']
+    out = [tf.ensure_shape(tf.io.parse_tensor(x[0], s.dtype), s.shape) for (x, s) in zip(
+        tf.split(out, len(tf.nest.flatten(spec, expand_composites=True))), 
+        tf.nest.flatten(spec, expand_composites=True))]
+    out = tf.nest.pack_sequence_as(spec, tf.nest.flatten(out), expand_composites=True)
+    return out
+
+def save_spec(path: str, spec: tensors.GraphTensor.Spec) -> None:
+    proto = spec.experimental_as_proto()
+    with open(path, 'wb') as fh:
+        fh.write(proto.SerializeToString())
+
+def load_spec(path: str) -> tensors.GraphTensor.Spec:
+    with open(path, 'rb') as fh:
+        serialized_proto = fh.read()
+    spec = tensors.GraphTensor.Spec.experimental_from_proto(
+        tensors.GraphTensor.Spec
+        .experimental_type_proto()
+        .FromString(serialized_proto))
+    return spec
+    

molcraft/tensors.py

@@ -0,0 +1,527 @@
+import tensorflow as tf 
+import keras
+import typing 
+from tensorflow.python.framework import composite_tensor
+
+from molcraft import ops
+
+
+class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
+    
+    def batch(self, spec: 'GraphTensor.Spec', batch_size: int | None):
+        def batch_field(f):
+            if isinstance(f, tf.TensorSpec):
+                return tf.TensorSpec(
+                    shape=[None] + f.shape[1:],
+                    dtype=f.dtype)
+            elif isinstance(f, tf.RaggedTensorSpec):
+                return tf.RaggedTensorSpec(
+                    shape=[batch_size, None] + f.shape[1:],
+                    dtype=f.dtype,
+                    ragged_rank=1,
+                    row_splits_dtype=f.row_splits_dtype)
+            elif isinstance(f, tf.TypeSpec):
+                return f.__batch_encoder__.batch(f, batch_size)
+            return f
+        fields = dict(spec.__dict__)
+        # Pop context fields as they will be batched differently
+        context_fields = fields.pop('context')
+        batched_fields = tf.nest.map_structure(batch_field, fields)
+        batched_spec = object.__new__(type(spec))
+        batched_context_fields = tf.nest.map_structure(
+            lambda spec: tf.TensorSpec([batch_size] + spec.shape, spec.dtype), 
+            context_fields)
+        batched_spec.__dict__.update({'context': batched_context_fields})
+        batched_spec.__dict__.update(batched_fields)
+        return batched_spec
+    
+    def unbatch(self, spec: 'GraphTensor.Spec'):
+        def unbatch_field(f):   
+            if isinstance(f, tf.TensorSpec):
+                return tf.TensorSpec(
+                    shape=[None] + f.shape[1:],
+                    dtype=f.dtype)
+            elif isinstance(f, tf.RaggedTensorSpec):
+                return tf.RaggedTensorSpec(
+                    shape=[None] + f.shape[2:],
+                    dtype=f.dtype,
+                    ragged_rank=0,
+                    row_splits_dtype=f.row_splits_dtype)
+            elif isinstance(f, tf.TypeSpec):
+                return f.__batch_encoder__.unbatch(f)
+            return f
+        fields = dict(spec.__dict__)
+        # Pop context fields as they will be unbatched differently
+        context_fields = fields.pop('context')
+        unbatched_fields = tf.nest.map_structure(unbatch_field, fields)
+        unbatched_context_fields = tf.nest.map_structure(
+            lambda spec: tf.TensorSpec(spec.shape[1:], spec.dtype), 
+            context_fields)
+        unbatched_spec = object.__new__(type(spec))
+        unbatched_spec.__dict__.update({'context': unbatched_context_fields})
+        unbatched_spec.__dict__.update(unbatched_fields)
+        return unbatched_spec
+        
+    def encode(self, spec: 'GraphTensor.Spec', value: 'GraphTensor', minimum_rank: int = 0):
+        unflatten = False if (is_ragged(spec) or is_scalar(spec)) else True 
+        if unflatten:
+            value = value.unflatten()
+        value_components = tuple(value.__dict__[key] for key in spec.__dict__)
+        value_components = tuple(
+            x for x in tf.nest.flatten(value_components) 
+            if isinstance(x, (tf.Tensor, composite_tensor.CompositeTensor))
+        )
+        return value_components
+    
+    def encoding_specs(self, spec: 'GraphTensor.Spec'):
+        def encode_fields(f):
+            if isinstance(f, tf.TensorSpec):
+                scalar = is_scalar(spec)
+                return tf.RaggedTensorSpec(
+                    shape=([None] if scalar else [None, None]) + f.shape[1:], 
+                    dtype=f.dtype, 
+                    ragged_rank=(0 if scalar else 1),
+                    row_splits_dtype=spec.context['size'].dtype)
+            return f
+        fields = dict(spec.__dict__)
+        context_fields = fields.pop('context')
+        encoded_fields = tf.nest.map_structure(encode_fields, fields)
+        encoded_fields = {**{'context': context_fields}, **encoded_fields}
+        spec_components = tuple(encoded_fields.values())
+        spec_components = tuple(
+            x for x in tf.nest.flatten(spec_components) 
+            if isinstance(x, tf.TypeSpec)
+        )
+        return spec_components
+    
+    def decode(self, spec, encoded_value):
+        spec_tuple = tuple(spec.__dict__.values())
+        encoded_value = iter(encoded_value)
+        value_tuple = [
+            next(encoded_value) if isinstance(x, tf.TypeSpec) else x
+            for x in tf.nest.flatten(spec_tuple)
+        ]
+        value_tuple = tf.nest.pack_sequence_as(spec_tuple, value_tuple)
+        fields = dict(zip(spec.__dict__.keys(), value_tuple))
+        value = object.__new__(spec.value_type)
+        value.__dict__.update(fields)
+
+        flatten = is_ragged(value) and not is_ragged(spec)
+        if flatten:
+            value = value.flatten()
+        return value
+    
+
+class GraphTensor(tf.experimental.BatchableExtensionType):
+    context: typing.Mapping[str, typing.Union[tf.Tensor, tf.RaggedTensor]]
+    node: typing.Mapping[str, typing.Union[tf.Tensor, tf.RaggedTensor]]
+    edge: typing.Mapping[str, typing.Union[tf.Tensor, tf.RaggedTensor]]
+
+    __batch_encoder__ = GraphTensorBatchEncoder()
+
+    __name__ = 'GraphTensor'
+
+    def __validate__(self):
+        assert 'size' in self.context, "graph.context['size'] is required."
+        assert self.context['size'].dtype == tf.int32, (
+            "dtype of graph.context['size'] needs to be int32.")
+        assert 'feature' in self.node, "graph.node['feature'] is required."
+        assert 'source' in self.edge, "graph.edge['source'] is required."
+        assert 'target' in self.edge, "graph.edge['target'] is required."
+        assert self.edge['source'].dtype == tf.int32, (
+            "dtype of graph.edge['source'] needs to be int32.")
+        assert self.edge['target'].dtype == tf.int32, (
+            "dtype of graph.edge['target'] needs to be int32.")
+        # TODO: Assert node sizes (based on context['size'])
+        #       Assert edge sizes (based on edge['source'])
+        
+    @property 
+    def spec(self):
+        def unspecify_spec(s):
+            if isinstance(s, tf.TensorSpec):
+                return tf.TensorSpec([None] + s.shape[1:], s.dtype)
+            return s
+        orig_spec = tf.type_spec_from_value(self)
+        fields = dict(orig_spec.__dict__)
+        context_fields = fields.pop('context')
+        new_spec_components = tf.nest.map_structure(unspecify_spec, fields)
+        new_spec_components['context'] = context_fields
+        return orig_spec.__class__(**new_spec_components)
+    
+    @property 
+    def shape(self):
+        if is_ragged(self):
+            return self.node['feature'].shape 
+        return self.context['size'].shape + [None] + self.node['feature'].shape[1:] 
+    
+    @property 
+    def dtype(self):
+        return self.node['feature'].dtype
+
+    @property
+    def graph_indicator(self):
+        dtype = self.context['size'].dtype
+        if is_scalar(self):
+            return tf.zeros(tf.shape(self.node['feature'])[:1], dtype=dtype)
+        num_graphs = keras.ops.shape(self.context['size'])[0]
+        if is_ragged(self):
+            return keras.ops.arange(num_graphs, dtype=dtype)
+        return keras.ops.repeat(keras.ops.arange(num_graphs, dtype=dtype), self.context['size'])
+
+    @property
+    def num_subgraphs(self) -> tf.Tensor:
+        dtype = self.context['size'].dtype
+        if is_scalar(self):
+            num_subgraphs = tf.constant(1, dtype=dtype)
+        else:
+            num_subgraphs = tf.shape(self.context['size'], out_type=dtype)[0]
+        if tf.executing_eagerly():
+            return int(num_subgraphs)
+        return num_subgraphs
+    
+    @property 
+    def num_nodes(self):
+        num_nodes = keras.ops.shape(self.node['feature'])[0]
+        if tf.executing_eagerly():
+            return int(num_nodes)
+        return num_nodes 
+    
+    @property 
+    def num_edges(self):
+        num_edges = keras.ops.shape(self.edge['source'])[0]
+        if tf.executing_eagerly():
+            return int(num_edges)
+        return num_edges 
+    
+    def gather(self, node_attr: str, edge_type: str) -> tf.Tensor:
+        if edge_type != 'source' and edge_type != 'target':
+            raise ValueError
+        return ops.gather(self.node[node_attr], self.edge[edge_type])
+    
+    def aggregate(self, edge_attr: str, edge_type: str = 'target') -> tf.Tensor:
+        if edge_type != 'source' and edge_type != 'target':
+            raise ValueError
+        edge_attr = self.edge[edge_attr]
+        if 'weight' in self.edge:
+            edge_attr = edge_attr * self.edge['weight']
+        return ops.aggregate(edge_attr, self.edge[edge_type], self.num_nodes)
+    
+    def propagate(self, add_edge_feature: bool = False):
+        updated_feature = ops.propagate(
+            node_feature=self.node['feature'],
+            edge_source=self.edge['source'],
+            edge_target=self.edge['target'],
+            edge_feature=self.edge.get('feature') if add_edge_feature else None,
+            edge_weight=self.edge.get('weight'),
+        )
+        return self.update({'node': {'feature': updated_feature}})
+
+    def flatten(self):
+        if not is_ragged(self):
+            raise ValueError(
+                f"{self.__class__.__qualname__} instance is already flat.")
+        def flatten_fn(x):
+            if isinstance(x, tf.RaggedTensor):
+                return x.flat_values
+            return x
+        edge_increment = ops.gather(
+            self.node['feature'].row_starts(), self.edge['source'].value_rowids())
+        edge_increment = tf.cast(
+            edge_increment, dtype=self.edge['source'].dtype)
+        data = to_dict(self)
+        flat_values = tf.nest.map_structure(flatten_fn, data)
+        flat_values['edge']['source'] += edge_increment
+        flat_values['edge']['target'] += edge_increment 
+        return from_dict(flat_values)
+        
+    def unflatten(self, *, force: bool = False):
+        if is_scalar(self):
+            raise ValueError(
+                f"{self.__class__.__qualname__} instance is a scalar " 
+                "and cannot be unflattened.")
+        if is_ragged(self):
+            raise ValueError(
+                f"{self.__class__.__qualname__} instance is already unflat.")
+        def unflatten_fn(x, value_rowids, nrows) -> tf.RaggedTensor:
+            if isinstance(x, tf.Tensor):
+                return tf.RaggedTensor.from_value_rowids(x, value_rowids, nrows)
+            return x
+        graph_indicator_node = self.graph_indicator
+        graph_indicator_edge = ops.gather(graph_indicator_node, self.edge['source'])
+        if force:
+            sorted_indices = keras.ops.argsort(graph_indicator_edge)
+        num_subgraphs = self.num_subgraphs
+        unflat_values = {}
+        data = to_dict(self)
+        for key in tf.type_spec_from_value(self).__dict__:
+            value = data[key]
+            if key == 'context':
+                unflat_values[key] = value
+            elif key == 'node':
+                unflat_values[key] = tf.nest.map_structure(
+                    lambda x: unflatten_fn(x, graph_indicator_node, num_subgraphs),
+                    value)
+                row_starts = unflat_values[key]['feature'].row_starts()
+                edge_decrement = ops.gather(row_starts, graph_indicator_edge)
+                if force:
+                    edge_decrement = ops.gather(edge_decrement, sorted_indices)
+                    graph_indicator_edge = ops.gather(graph_indicator_edge, sorted_indices)
+                edge_decrement = tf.cast(edge_decrement, dtype=self.edge['source'].dtype)
+            elif key == 'edge':
+                if force:
+                    value = tf.nest.map_structure(lambda x: ops.gather(x, sorted_indices), value)
+                value['source'] -= edge_decrement
+                value['target'] -= edge_decrement
+                unflat_values[key] = tf.nest.map_structure(
+                    lambda x: unflatten_fn(x, graph_indicator_edge, num_subgraphs),
+                    value)
+        return from_dict(unflat_values)
+
+    def update(self, values):
+        data = to_dict(self)
+        for outer_field, mapping in values.items():
+            if outer_field == 'context':
+                reference_value = data[outer_field]['size']
+            elif outer_field == 'edge':
+                reference_value = data[outer_field]['source']
+            else:
+                reference_value = data[outer_field]['feature']
+            for inner_field, value in mapping.items():
+                if value is None:
+                    data[outer_field].pop(inner_field, None)
+                    continue
+                data[outer_field][inner_field] = _maybe_convert_new_value(
+                    value, reference_value
+                )
+        return self.__class__(**data)
+
+    def replace(self, values):
+        data = to_dict(self)
+        for outer_field, mapping in values.items():
+            if outer_field == 'context':
+                reference_value = data[outer_field]['size']
+            elif outer_field == 'edge':
+                reference_value = data[outer_field]['source']
+            else:
+                reference_value = data[outer_field]['feature']
+            for inner_field, value in mapping.items():
+                values[outer_field][inner_field] = _maybe_convert_new_value(
+                    value, reference_value
+                )
+        return self.__class__(**values)
+    
+    def __getitem__(self, index):
+        if index is None and is_scalar(self):
+            return self.__class__(
+                context={key: value[None] for (key, value) in self.context.items()},
+                node=self.node,
+                edge=self.edge,
+            )
+        if not is_ragged(self):
+            is_flat = True
+            tensor = self.unflatten()
+        else:
+            tensor = self
+            is_flat = False
+        data = to_dict(tensor)
+        if isinstance(index, (slice, int)):
+            data = tf.nest.map_structure(lambda x: x[index], data)
+        else:
+            data = tf.nest.map_structure(lambda x: ops.gather(x, index), data)
+        tensor = from_dict(data)
+        if is_flat and not is_scalar(tensor):
+            return tensor.flatten()
+        return tensor
+    
+    def __repr__(self):
+        return _repr(self)
+    
+    def numpy(self):
+        """For now added to work with `keras.Model.predict`"""
+        return self 
+
+    class Spec:
+
+        def __init__(
+            self, 
+            context: typing.Mapping[str, tf.TensorSpec | tf.RaggedTensorSpec], 
+            node: typing.Mapping[str, tf.TensorSpec | tf.RaggedTensorSpec],  
+            edge: typing.Mapping[str, tf.TensorSpec | tf.RaggedTensorSpec],  
+        ) -> None:
+            self.context = context 
+            self.node = node
+            self.edge = edge 
+
+        @property
+        def shape(self):
+            if is_ragged(self):
+                return self.node['feature'].shape 
+            return self.context['size'].shape + [None] + self.node['feature'].shape[1:] 
+
+        @classmethod
+        def from_input_shape_dict(cls, input_shape: dict[str, tf.TensorShape]) -> 'GraphTensor.Spec':
+            for key, value in input_shape.items():
+                input_shape[key] = {k: tf.TensorShape(v) for k, v in value.items()}
+            return cls(**tf.nest.map_structure(lambda s: tf.TensorSpec(s, dtype=tf.variant), input_shape)) 
+
+        def __repr__(self):
+            return _repr(self)
+    
+    
+@tf.experimental.dispatch_for_api(tf.concat, {'values': typing.List[GraphTensor]})
+def graph_tensor_concat(
+    values: typing.List[GraphTensor],
+    axis: int = 0,
+    name: str = 'concat'
+) -> GraphTensor:
+    ragged = [is_ragged(v) for v in values]
+    if 0 < sum(ragged) < len(ragged):
+        raise ValueError(
+            'Nested data of the GraphTensor instances do not have consistent '
+            'types: found both tf.RaggedTensor values and tf.Tensor values.')
+    else:
+        ragged = ragged[0]
+
+    if ragged:
+        values = [v.flatten() for v in values]
+
+    flat_values = [tf.nest.flatten(v, expand_composites=True) for v in values]
+    flat_values = [tf.concat(f, axis=0) for f in list(zip(*flat_values))]
+    num_edges = [keras.ops.shape(v.edge['source'])[0] for v in values]
+    num_nodes = [keras.ops.shape(v.node['feature'])[0] for v in values]
+    incr = tf.concat([[0], tf.cumsum(num_nodes)[:-1]], axis=0)
+    incr = tf.repeat(incr, num_edges)
+    value = tf.nest.pack_sequence_as(values[0], flat_values, expand_composites=True)
+
+    edge_update = {
+        'source': value.edge['source'] + incr,
+        'target': value.edge['target'] + incr,
+    }
+    value = value.update(
+        {
+            'edge': edge_update
+        },
+    )
+    if not ragged:
+        return value
+    return value.unflatten()
+
+# TODO: Clean this up.
+@tf.experimental.dispatch_for_api(tf.stack, {'values': typing.List[GraphTensor]})
+def graph_tensor_stack(
+    values: typing.List[GraphTensor],
+    axis: int = 0,
+    name: str = 'stack'
+) -> GraphTensor:
+    ragged = [is_ragged(v) for v in values]
+    if not is_scalar(values[0]):
+        raise ValueError(
+            'tf.stack on a list of `GraphTensor`s is currently '
+            'only supported for scalar `GraphTensor`s. '
+        )
+    if any(ragged):
+        raise ValueError(
+            'tf.stack on a list of `GraphTensor`s is currently '
+            'only supported for flattened `GraphTensor`s. '
+        )
+    
+    def concat_or_stack(k, v):
+        if k.startswith('context'):
+            return tf.stack(v, axis=0)
+        return tf.concat(v, axis=0)
+
+    fields = tuple(tf.type_spec_from_value(values[0]).__dict__)
+    num_inner_fields = tuple(len(values[0].__dict__[field]) for field in fields)
+    outer_keys = []
+    for (f, num_fields) in zip(fields, num_inner_fields):
+        outer_keys.extend([f] * num_fields)
+
+    flat_values = [tf.nest.flatten(v, expand_composites=True) for v in values]
+    flat_values = [concat_or_stack(k, f) for k, f in zip(outer_keys, list(zip(*flat_values)))]
+    value = tf.nest.pack_sequence_as(values[0], flat_values, expand_composites=True)
+
+    num_edges = [keras.ops.shape(v.edge['source'])[0] for v in values]
+    num_nodes = [keras.ops.shape(v.node['feature'])[0] for v in values]
+    incr = tf.concat([[0], tf.cumsum(num_nodes)[:-1]], axis=0)
+    incr = tf.repeat(incr, num_edges)
+    edge_update = {
+        'source': value.edge['source'] + incr,
+        'target': value.edge['target'] + incr,
+    }
+    value = value.update(
+        {
+            'edge': edge_update
+        },
+    )
+    return value
+
+def is_scalar(value_or_spec: GraphTensor | GraphTensor.Spec) -> bool:
+    return value_or_spec.context['size'].shape.rank == 0
+
+def is_ragged(value_or_spec: GraphTensor | GraphTensor.Spec) -> bool:
+    is_ragged = isinstance(
+        value_or_spec.node['feature'], (tf.RaggedTensor, tf.RaggedTensorSpec))
+    if isinstance(value_or_spec, tf.RaggedTensorSpec):
+        is_ragged = (
+            is_ragged and value_or_spec.node['feature'].ragged_rank == 1)
+    return is_ragged
+
+def to_dict(tensor: GraphTensor) -> dict:
+    spec = tf.type_spec_from_value(tensor)
+    return {key: dict(tensor.__dict__[key]) for key in spec.__dict__}
+
+def from_dict(data: dict) -> GraphTensor:
+    data['context']['size'] = tf.cast(data['context']['size'], tf.int32)
+    data['edge']['source'] = tf.cast(data['edge']['source'], tf.int32)
+    data['edge']['target'] = tf.cast(data['edge']['target'], tf.int32)
+    return GraphTensor(**data)
+
+def is_graph(data):
+    if isinstance(data, GraphTensor):
+        return True 
+    elif isinstance(data, dict) and 'size' in data.get('context', {}):
+        return True 
+    return False
+    
+def _maybe_convert_new_value(
+    new_value: tf.Tensor | tf.RaggedTensor, 
+    old_value: tf.Tensor | tf.RaggedTensor | None,
+) -> tf.Tensor | tf.RaggedTensor:
+    if old_value is None:
+        return new_value
+    is_old_ragged = isinstance(old_value, tf.RaggedTensor)
+    is_new_ragged = isinstance(new_value, tf.RaggedTensor)
+    if is_old_ragged and not is_new_ragged:
+        new_value = old_value.with_flat_values(new_value)
+    elif not is_old_ragged and is_new_ragged:
+        new_value = new_value.flat_values
+    return new_value
+
+def _repr(x: GraphTensor | GraphTensor.Spec):
+    if isinstance(x, GraphTensor):
+        def _trepr(v: tf.Tensor | tf.RaggedTensor):
+            if isinstance(v, tf.Tensor):
+                return f'<tf.Tensor: shape={v.shape.as_list()}, dtype={v.dtype.name}>'
+            return (
+                f'<tf.RaggedTensor: shape={v.shape.as_list()}, ' 
+                f'dtype={v.dtype.name}, ragged_rank={v.ragged_rank}>'
+            )
+    else:
+        def _trepr(v: tf.TensorSpec | tf.RaggedTensorSpec):
+            if isinstance(v, tf.TensorSpec):
+                return f'<tf.TensorSpec: shape={v.shape.as_list()}, dtype={v.dtype.name}>'
+            return (
+                f'<tf.RaggedTensorSpec: shape={v.shape.as_list()}, ' 
+                f'dtype={v.dtype.name}, ragged_rank={v.ragged_rank}>'
+            )
+        
+    context_fields = f',\n        '.join([f'{k!r}: {_trepr(v)}' for k, v in x.context.items()])
+    node_fields = f',\n        '.join([f'{k!r}: {_trepr(v)}'for k, v in x.node.items()])
+    edge_fields = f',\n        '.join([f'{k!r}: {_trepr(v)}' for k, v in x.edge.items()])
+
+    context_field = 'context={\n        ' + context_fields + '\n    }' 
+    node_field = 'node={\n        ' + node_fields + '\n    }' 
+    edge_field = 'edge={\n        ' + edge_fields + '\n    }' 
+
+    fields = ',\n    '.join([context_field, node_field, edge_field])
+    return x.__class__.__name__ + '(\n    ' + fields + '\n)'