Terrace API

terrace.batch module

class terrace.batch.Batch(items: List[T] | Type[T] | None, **kwargs)

Bases: BatchBase[T]

asdict()

Convert to dict

attribute_names() → List[str]

Returns the names of all the batched attributes

item_type() → Type[T]

Returns the type of each item in the batch

to(device)

class terrace.batch.BatchBase(*args, **kwds)

Bases: Generic[T]

cpu()

cuda()

item_type() → Type[Batchable]

Returns the type of each item in the batch

class terrace.batch.BatchView(batch: Batch, index: int | slice)

Bases: BatchViewBase[T]

View of an item in a batch. Should act like said item in most circumstances. We use views instead of creating actual items because, for many use cases, lazily indexing batches is much faster

asdict()

Override this method to define what attributes you want your Batch to define

get_type()

class terrace.batch.BatchViewBase(*args, **kwds)

Bases: Generic[T], Batchable

class terrace.batch.Batchable

Bases: object

Base class for all objects we want to batchify

This method can also define static methods collate_{attribute} and index_{attribute}

asdict()

Override this method to define what attributes you want your Batch to define

static get_batch_type()

override this method if you want a custom batch type for your batchable class

class terrace.batch.DataLoader(dataset, batch_size=1, shuffle=False, **kwargs)

Bases: DataLoader

Dataloader that correctly batchifies Batchable data

batch_size: int | None

dataset: Dataset[T_co]

drop_last: bool

num_workers: int

pin_memory: bool

pin_memory_device: str

prefetch_factor: int | None

sampler: Sampler | Iterable

timeout: float

class terrace.batch.LazyBatch(items)

Bases: BatchBase[T]

This is mainly used if you’re recollating after indexing a batch ( eg collate([batch[0]])). If you use collate([batch[0]], lazy=True) (recommended), it will return a LazyBatch. This is nice because it will only re-collate the attributes of the batch on the fly

to(device)

class terrace.batch.NoStackTensor(tensor: Tensor)

Bases: Tensor

This is used when you want to collate tensors into a list instead of stack them. E.g. when you have different shapes

terrace.batch.collate(batch: Any, lazy=False) → Any

turn a list of items into a batch of items. Replacement for pytorch’s default collate. This is what we use in the custom DataLoader class

terrace.categorical_tensor module

class terrace.categorical_tensor.CategoricalTensor(tensor: Tensor, num_classes: int | Tuple[int, ...])

Bases: Tensor

Subclass of torch Tensors with dtype long. They have an additional num_classes member that is either an int of a tuple. If num_classes is an int, it is the number of classes in the tensor as whole (that is, all numbers in the tensor are in the range [0, num_classes)), If num_classes is a tuple, is is the number of classes along the last dimension of the tensor. For instance, if t.num_classes is (4,8), then t[...0] has 4 classes and t[...,1] has 8 classes.

class terrace.categorical_tensor.NoStackCatTensor(tensor: Tensor, num_classes: int | Tuple[int, ...])

Bases: CategoricalTensor

terrace.dataframe module

class terrace.dataframe.DFRow(**kwargs)

Bases: Batchable

items()

keys()

values()

terrace.dataframe.merge(items: DFRow | Batch)

terrace.graph module

class terrace.graph.Graph(nodes: Sequence[N], edges: Sequence[Tuple[int, int]], edata: List[E] | None = None, directed: bool = False, device='cpu')

Bases: GraphBase[N, E], Batchable

Wrapper around dgl graph allowing easier access to data

static get_batch_type()

override this method if you want a custom batch type for your batchable class

class terrace.graph.GraphBase(_dgl_graph: graph | batch, _node_type_tree: TypeTree, _edge_type_tree: TypeTree | None)

Bases: Generic[N, E]

dgl() → graph | batch

property edata: Batch[E]

property edges: List[Tuple[int, int]]

property ndata: Batch[N]

to(device)

class terrace.graph.GraphBatch(items: List[Graph[N, E]])

Bases: BatchBase[G], GraphBase

edge_slices: List[slice]

item_type()

Returns the type of each item in the batch

node_slices: List[slice]

class terrace.graph.GraphBatchView(batch: Batch, index: int | slice)

Bases: BatchViewBase[G]

dgl()

property edata: Batch[E]

property edges: List[Tuple[int, int]]

static get_batch_type()

override this method if you want a custom batch type for your batchable class

get_type()

property ndata: Batch[N]

class terrace.graph.TypeTree(type: Type, subtypes: Dict[str, TypeTree])

Bases: object

subtypes: Dict[str, TypeTree]

type: Type

terrace.graph.flatten_batch(b: Batch) → Tuple[Dict[str, Tensor], TypeTree]

Flattens batch to a single dict from keys to tensors. Also returns all the batch types. This is neccessary because dgl stores ndata and edata as such dicts.

terrace.graph.make_batch_from_unflat_dict(unflat_dict, type_tree)

Helper function for unflatten_batch. Makes a batch from a dict that has been unflattened

terrace.graph.unflatten_batch(flat_dict: Dict[str, Tensor], type_tree: TypeTree) → Batch

Reverses flatten_batch; takes a flattened dict and a type tree and returns a proper Batch

terrace.graph.unflatten_dict(flat)

Helper function for unflatten_batch. Returns an unflattened dict for a dict with keys like ‘key/subkey’

terrace.module module

class terrace.module.LazyEmbedding(embedding_dims: Tuple[int] | int)

Bases: Module

LazyEmbedding uses the num_classes from CategoricalTensors to determine embedding weight size. Note that it assumes tensors have shape (…, N) where N is the number of categorical features. So, in most cases where you have a batch of single categorical features, you must give the embedding a tensor of shape (B, 1). Admittedly this is a bit weird, but it does nicely extend to cases where you have multiple categorical features. In this case, the embedding creates an nn.Embedding layer for each feature and concatenates the result together. Thus the output of this layer has shape (B, E*N), where E is the embedding_dim.

forward(x: CategoricalTensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class terrace.module.LazyLayerNorm(*args, **kwargs)

Bases: WrapperModule

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class terrace.module.LazyLinear(*args, **kwargs)

Bases: WrapperModule

torch >=1.13 has this, but wanted to try my hand at implementing it myself (and ensuring that older torch versions work)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class terrace.module.LazyMultiheadAttention(*args, **kwargs)

Bases: WrapperModule

forward(q, k, v)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class terrace.module.Module

Bases: Module

checkpoint()

is_initialized()

loop_body()

loop_start()

make(cls, *args, **kwargs)

make_param(cls, *args, **kwargs)

parameters(recurse: bool = True)

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

start_forward()

training: bool

class terrace.module.WrapperModule(*args, **kwargs)

Bases: Module

training: bool