tokitai-operator 0.1.0

Verified DL kernel compiler: formally-checked GEMM, p-adic, sheaf, contract-carrying ops. Paper-artifact grade.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

//! Adapter from [`LocalSample`] to the [`Tensor<f32>`] shapes expected
//! by the model layer / training driver.
//!
//! ## Why an adapter at all?
//!
//! The IR [`Tensor<f32>`] type carries a [`Shape`] and a
//! [`DomainId`]. The training driver
//! (`crates/training/src/step.rs::train_step`) calls the user-supplied
//! `forward` / `backward` closures with `&[Tensor<f32>]` and
//! `&Tensor<f32>`. The MoE model's first layer expects:
//!
//! - Input: `Tensor<f32>` of shape `[1, 96]` (batch=1, feature_dim=96).
//! - Target: `Tensor<f32>` of shape `[1, 20]` (batch=1, label_dim=20).
//!
//! (Phase 2.5 uses batch=1 because the existing `train_step` driver
//! is single-sample. Batched iteration is a follow-up phase that
//! stacks N `[1, 96]` tensors into `[N, 96]` via `Tensor::concat`.)
//!
//! ## Encoding
//!
//! `LocalSample.features` is a flat 96-dim `Vec<f32>` already; we
//! just wrap it with the right shape. Same for `labels`.
//!
//! - [`to_input_tensor`] returns `Tensor<f32>` with shape `[1, 96]`
//!   and the feature data laid out in **row-major** (C) order: the
//!   first 74 dims are the categorical one-hot (cast to `f32` 0.0/1.0)
//!   and the last 22 dims are the numerical milli-units. The
//!   categorical one-hots use values in `{0.0, 1.0}` so the
//!   [`Embedding`] / [`Router`] layers can index them directly.
//! - [`to_target_tensor`] returns `Tensor<f32>` with shape `[1, 20]`
//!   and the label data laid out in row-major order: the first 12
//!   dims are the outcome one-hot (0.0/1.0) and the last 8 dims are
//!   the aux metric scalars in `[0, 1]`.

use crate::dataset_bridge::local_dataset::LocalSample;
use crate::dataset_bridge::{FEATURE_DIM, LABEL_DIM};
use crate::domain::DomainId;
use crate::error::{Error, Result};
use crate::object::{Dim, Shape, Tensor};

/// Convert a [`LocalSample`] into a `[1, 96]` input tensor.
pub fn to_input_tensor(sample: &LocalSample) -> Result<Tensor<f32>> {
    if sample.features.len() != FEATURE_DIM {
        return Err(Error::shape(format!(
            "to_input_tensor: expected features.len() == {}, got {}",
            FEATURE_DIM,
            sample.features.len()
        )));
    }
    Ok(Tensor::dense_cpu(
        DomainId::new("f32"),
        Shape::new(vec![Dim::Static(1), Dim::Static(FEATURE_DIM)]),
        sample.features.clone(),
    ))
}

/// Convert a [`LocalSample`] into a `[1, 20]` target tensor.
pub fn to_target_tensor(sample: &LocalSample) -> Result<Tensor<f32>> {
    if sample.labels.len() != LABEL_DIM {
        return Err(Error::shape(format!(
            "to_target_tensor: expected labels.len() == {}, got {}",
            LABEL_DIM,
            sample.labels.len()
        )));
    }
    Ok(Tensor::dense_cpu(
        DomainId::new("f32"),
        Shape::new(vec![Dim::Static(1), Dim::Static(LABEL_DIM)]),
        sample.labels.clone(),
    ))
}