Tensor Datasets

Conversion between leading-axis tensors, tensor packs, and runtime datasets.

TensorDataset (leading-axis batching)

PyTorch's TensorDataset concept is: given one or more tensors that share the same size(0), build a dataset of samples by slicing each tensor along its leading batch axis.

In TorchLean we do the same thing, but with shapes tracked in the type:

• a batched tensor has shape .dim n σ,
• slicing at i : Fin n yields a sample of shape σ,
• and a batch of multiple tensors is represented as a TensorPack.

def TorchLean.Data.unbatchTensorPack {β : Type} {n : ℕ} {ss : List Spec.Shape} : TensorPack β (List.map (fun (s : Spec.Shape) => Spec.Shape.dim n s) ss) → Fin n → TensorPack β ss

Slice a batched TensorPack along its leading batch axis.

If a sample is represented as a shape-indexed tuple TensorPack β ss, then a minibatch of size n is TensorPack β (ss.map (fun s => .dim n s)). This function picks a batch index i : Fin n and returns the corresponding single sample.

def TorchLean.Data.castTListOfFloat {α : Type} [Runtime.FromFloat α] {ss : List Spec.Shape} : TensorPack Float ss → TensorPack α ss

Convert a shape-indexed TensorPack of Float tensors to the runtime scalar type α.

def TorchLean.Data.tensorDatasetFromLeadingAxis {β : Type} {n : ℕ} {ss : List Spec.Shape} (xs : TensorPack β (List.map (fun (s : Spec.Shape) => Spec.Shape.dim n s) ss)) : Dataset (TensorPack β ss)

Build a dataset by slicing a batched TensorPack along the leading batch axis. This gives the typed counterpart of a tensor dataset built from several aligned arrays.

def TorchLean.Data.tensorDatasetFromLeadingAxisFloat {α : Type} [Runtime.FromFloat α] {n : ℕ} {ss : List Spec.Shape} (xs : TensorPack Float (List.map (fun (s : Spec.Shape) => Spec.Shape.dim n s) ss)) : Dataset (TensorPack α ss)

Float-to-α variant of tensorDatasetFromLeadingAxis, for data loaded from disk.

def TorchLean.Data.supervisedFromLeadingAxis {α : Type} {n : ℕ} {σ τ : Spec.Shape} (X : Spec.Tensor α (Spec.Shape.dim n σ)) (Y : Spec.Tensor α (Spec.Shape.dim n τ)) : Dataset (TensorPack α [σ, τ])

Supervised dataset from two batched tensors X : (n, σ) and Y : (n, τ) by slicing the leading batch axis.

This is the common regression/supervised-learning case: the TorchLean analogue of TensorDataset(X, Y) in PyTorch.

def TorchLean.Data.supervisedFromLeadingAxisFloat {α : Type} [Runtime.FromFloat α] {n : ℕ} {σ τ : Spec.Shape} (X : Spec.Tensor Float (Spec.Shape.dim n σ)) (Y : Spec.Tensor Float (Spec.Shape.dim n τ)) : Dataset (TensorPack α [σ, τ])

Float-to-α variant of supervisedFromLeadingAxis, for data loaded from disk.