PyTorch → flodl Migration Guide

use flodl::*;  // brings in Tensor, Variable, all nn modules, graph builder, etc.

use flodl::tensor::{Tensor, TensorOptions, DType, Device};
use flodl::autograd::{Variable, no_grad};
use flodl::nn::{Linear, Adam, mse_loss, Module, Parameter};
use flodl::graph::{FlowBuilder, Graph, MergeOp};

let y = x.matmul(&w)?.add(&b)?.relu()?;

// flodl
flodl::manual_seed(42);
flodl::cuda_manual_seed_all(42);  // no-op without CUDA feature

use flodl::Rng;

let mut rng = Rng::seed(42);       // deterministic
rng.shuffle(&mut indices);          // Fisher-Yates shuffle
let val = rng.f32();               // uniform [0, 1)
let coin = rng.bernoulli(0.5);    // true ~50%

// flodl
let opts = TensorOptions::default();  // Float32, CPU
let x = Tensor::zeros(&[2, 3], opts)?;
let x = Tensor::ones(&[2, 3], opts)?;
let x = Tensor::rand(&[2, 3], opts)?;
let x = Tensor::randn(&[2, 3], opts)?;
let x = Tensor::full(&[2, 3], 7.0, opts)?;
let x = Tensor::eye(4, opts)?;
let x = Tensor::arange(0.0, 10.0, 2.0, opts)?;
let x = Tensor::from_f32(&[1.0, 2.0, 3.0], &[3], Device::CPU)?;
let x = Tensor::from_i64(&[0, 1, 2], &[3], Device::CPU)?;
let x = Tensor::linspace(0.0, 1.0, 10, opts)?;

// Like-tensors (same dtype/device as source)
let y = Tensor::zeros_like(&x)?;
let y = Tensor::ones_like(&x)?;

// Stacking
let y = Tensor::stack(&[&a, &b, &c], 0)?;

// flodl
let c = a.add(&b)?;
let c = a.sub(&b)?;
let c = a.mul(&b)?;
let c = a.div(&b)?;
let c = a.matmul(&b)?;
let c = x.mul_scalar(2.0)?;
let c = x.add_scalar(1.0)?;
let c = x.div_scalar(3.0)?;
let c = x.neg()?;

// flodl
let y = x.exp()?;
let y = x.log()?;
let y = x.sqrt()?;
let y = x.abs()?;
let y = x.pow_scalar(2.0)?;
let y = x.clamp(-1.0, 1.0)?;
let y = x.clamp_min(0.0)?;
let y = x.clamp_max(1.0)?;
let y = x.sin()?;
let y = x.cos()?;
let y = x.tan()?;
let y = x.asin()?;
let y = x.acos()?;
let y = x.atan()?;
let y = y_var.atan2(&x_var)?;            // Variable method
let y = x.sign()?;
let y = x.floor()?;
let y = x.ceil()?;
let y = x.round()?;
let y = x.trunc()?;
let y = x.frac()?;
let y = x.reciprocal()?;
let y = x.log1p()?;                      // ln(1+x), stable for small x
let y = x.expm1()?;                      // exp(x)-1, stable for small x
let y = x.log2()?;
let y = x.log10()?;
let y = x.erf()?;
let y = x.erfc()?;
let y = x.fmod(3.0)?;                    // C-style remainder
let y = x.remainder(3.0)?;               // Python-style modulo
let y = start.lerp(&end, 0.5)?;          // linear interpolation
let y = bias.addmm(&mat1, &mat2, 1.0, 1.0)?;  // beta*self + alpha*(mat1 @ mat2)
let y = inp.addcmul(&t1, &t2, 1.0)?;    // self + value * t1 * t2
let y = inp.addcdiv(&t1, &t2, 1.0)?;    // self + value * t1 / t2
let y = x.isclose(&y, 1e-5, 1e-8)?;
let y = a.maximum(&b)?;
let y = a.minimum(&b)?;
let y = x.masked_fill(&mask, 0.0)?;      // Variable method
let y = x.normalize(2.0, 1)?;            // Lp-normalize along dim
let y = a.cosine_similarity(&b, 1, 1e-8)?;

// flodl
let y = x.relu()?;
let y = x.sigmoid()?;
let y = x.tanh()?;
let y = x.gelu()?;
let y = x.silu()?;
let y = x.leaky_relu(0.01)?;
let y = x.elu(1.0)?;
let y = x.softplus(1.0, 20.0)?;
let y = x.mish()?;
let y = x.selu()?;
let y = x.hardswish()?;
let y = x.hardsigmoid()?;
let y = x.softmax(1)?;
let y = x.log_softmax(1)?;

// flodl
let s = x.sum()?;
let s = x.sum_dim(1, true)?;
let m = x.mean()?;
let m = x.mean_dim(1, true)?;
let v = x.var()?;
let v = x.std()?;
let v = x.max_dim(1, true)?;
let v = x.min_dim(1, true)?;
let idx = x.argmax(1, false)?;
let n = x.norm()?;
let p = x.prod()?;
let p = x.prod_dim(1, true)?;
let c = x.cumsum(0)?;
let l = x.logsumexp(1, true)?;

// flodl
let y = x.reshape(&[2, 3])?;
// no separate view — reshape handles it
let y = x.squeeze(0)?;
let y = x.unsqueeze(0)?;
let y = x.flatten(1, -1)?;
let y = x.permute(&[0, 2, 1])?;
let y = x.transpose(0, 1)?;
let y = x.expand(&[4, 3])?;
let y = x.contiguous()?;
let y = x.movedim(0, 2)?;
let y = x.flip(&[0, 1])?;
let y = x.roll(2, 0)?;
let y = x.diagonal(0, 0, 1)?;
let y = x.tile(&[2, 3])?;
let y = x.triu(0)?;
let y = x.tril(0)?;
let parts = x.split(2, 0)?;
let slices = x.unbind(0)?;
let grids = Tensor::meshgrid(&[&x, &y])?;

// flodl
let y = x.select(0, 0)?;
let y = x.narrow(1, 1, 2)?;
let y = x.index_select(0, &indices)?;
let y = a.cat(&b, 0)?;
let y = Tensor::cat_many(&[&a, &b, &c], 0)?;  // concatenate many tensors
let y = Tensor::stack(&[&a, &b], 0)?;
let chunks = x.chunk(3, 1)?;
let y = x.repeat(&[2, 3])?;
let y = x.pad(&[1, 1], 0.0)?;                  // constant-value pad
let y = x.pad_mode(&[1, 1], 1, 0.0)?;          // 0=constant, 1=reflect, 2=replicate, 3=circular

// Batch iteration (split along dim 0)
for batch in data.batches(32)? {
    let x = Variable::new(batch, false);
    // ...
}

// flodl
let mask = x.gt(&threshold)?;
let mask = x.ge_scalar(0.0)?;
let y = Tensor::where_cond(&mask, &a, &b)?;

// Scalar comparisons
let mask = x.gt_scalar(0.0)?;
let mask = x.lt_scalar(1.0)?;
let mask = x.ge_scalar(0.0)?;
let mask = x.le_scalar(1.0)?;

// flodl
let y = x.to_dtype(DType::Float32)?;
let y = x.to_dtype(DType::Float64)?;
let y = x.to_dtype(DType::Float16)?;
let y = x.to_dtype(DType::BFloat16)?;

// flodl
let val = loss.item()?;           // scalar → f64
let data = x.to_f32_vec()?;      // → Vec<f32>
let data = x.to_f64_vec()?;      // → Vec<f64>
let data = x.to_i64_vec()?;      // → Vec<i64>

// flodl
x.shape()        // Vec<i64>: [2, 3]
x.ndim()         // usize: 2
x.numel()        // i64: 6
x.dtype()        // DType::Float32
x.device()       // Device::CPU

// flodl
let xt = Tensor::from_f32(&[1.0, 2.0], &[2], Device::CPU)?;
let x = Variable::new(xt, true);  // requires_grad = true
let y = x.pow_scalar(2.0)?.sum()?;
y.backward()?;
println!("{:?}", x.grad());  // Some(tensor([2.0, 4.0]))

// flodl
let layer = Linear::new(784, 128)?;
let layer = Linear::no_bias(784, 128)?;
let layer = Conv1d::configure(3, 16, 5).with_stride(2).with_padding(2).done()?;
let layer = Conv2d::new(3, 64, 3)?;                                          // defaults: stride=1, padding=0
let layer = Conv2d::configure(3, 64, 3).with_padding(1).with_stride(2).done()?;    // fluent builder
let layer = Conv2d::build(3, 64, 3, true, [1,1], [1,1], [1,1], 1, Device::CPU)?;   // full control
let layer = Conv3d::configure(1, 32, [3,3,3]).with_padding([1,1,1]).done()?;
let layer = ConvTranspose1d::new(16, 3, 5)?;
let layer = ConvTranspose2d::new(64, 3, 4)?;
let layer = ConvTranspose3d::new(32, 1, [3,3,3])?;
let pool = MaxPool2d::new(2);                                                // kernel=2, stride=2 (defaults to kernel)
let pool = MaxPool2d::with_stride(3, 2).padding(1);                          // kernel=3, stride=2, padding=1
let pool = AvgPool2d::new(2);
let pool = MaxPool1d::new(2);
let pool = AvgPool1d::new(2);
let pool = AdaptiveMaxPool2d::new(7, 7);
let pool = AdaptiveAvgPool2d::new([1, 1]);                                   // global avg pool (ResNet head)
let layer = PixelShuffle::new(2);
let layer = PixelUnshuffle::new(2);
let layer = Upsample::new(&[64, 64], 1);                                    // mode: 0=nearest, 1=bilinear
let layer = Unfold::new([3,3], [1,1], [0,0], [1,1]);
let layer = Fold::new([28,28], [3,3], [1,1], [0,0], [1,1]);
let layer = LayerNorm::new(128)?;
let layer = RMSNorm::new(128)?;
let layer = GroupNorm::new(4, 16)?;
let layer = BatchNorm::new(128)?;                                            // for [B, features] after Linear
let layer = BatchNorm2d::new(64)?;                                           // for [B, C, H, W] after Conv2d
let layer = InstanceNorm::new(64, true)?;                                    // affine=true
let layer = Dropout::new(0.5);
let layer = Dropout2d::new(0.1);
let layer = AlphaDropout::new(0.1);                                          // for SELU networks
let layer = ZeroPad2d::new(1);
let layer = ReflectionPad2d::new(1);
let layer = Embedding::new(1000, 128)?;
let layer = EmbeddingBag::new(1000, 128)?;
let cell = GRUCell::new(128, 256)?;
let cell = LSTMCell::new(128, 256)?;
let layer = GRU::new(128, 256, 2)?;                                          // 2-layer GRU
let layer = LSTM::new(128, 256, 2)?;                                         // 2-layer LSTM
let layer = MultiheadAttention::new(512, 8)?;
let layer = Bilinear::new(128, 64, 32, true)?;

// On a specific device (all modules have on_device() variants):
let layer = Linear::on_device(784, 128, Device::CUDA(0))?;
let layer = Conv2d::configure(3, 64, 3).with_padding(1).on_device(Device::CUDA(0)).done()?;
let layer = LayerNorm::on_device(128, Device::CUDA(0))?;
let layer = RMSNorm::on_device(128, Device::CUDA(0))?;
let layer = GroupNorm::on_device(4, 16, Device::CUDA(0))?;
let layer = BatchNorm::on_device(128, Device::CUDA(0))?;
let layer = BatchNorm2d::on_device(64, Device::CUDA(0))?;
let layer = InstanceNorm::on_device(64, true, Device::CUDA(0))?;
let layer = Embedding::on_device(1000, 128, Device::CUDA(0))?;
let cell = GRUCell::on_device(128, 256, Device::CUDA(0))?;
let cell = LSTMCell::on_device(128, 256, Device::CUDA(0))?;
let layer = GRU::on_device(128, 256, 2, false, Device::CUDA(0))?;
let layer = LSTM::on_device(128, 256, 2, false, Device::CUDA(0))?;
let layer = MultiheadAttention::on_device(512, 8, Device::CUDA(0))?;

// flodl — unit structs (no allocation, bare-name usage)
ReLU
Sigmoid
Tanh
GELU                       // erf form (default) — GELU::tanh() for the tanh approximation
SiLU
Mish
SELU
Hardswish
Hardsigmoid
Identity

// Parameterized at construction
LeakyReLU::new(0.01)
ELU::new(1.0)
Softplus::new(1.0, 20.0)         // beta, threshold
Softmax::new(-1)                  // dim
LogSoftmax::new(-1)
Flatten::new(1, -1)               // start_dim, end_dim

// Learnable parameters
let prelu = PReLU::new(1, Device::CPU)?;

// flodl — as a Module (for use in FlowBuilder graphs)
let blur = GaussianBlur::new(1.5);  // kernel size auto-computed from sigma
let y = blur.forward(&x)?;

// flodl — as a free function
let y = gaussian_blur_2d(&x, 1.5)?;  // input must be [B, C, H, W]

// flodl — declare children via parameters()
struct MLP {
    fc1: Linear,
    fc2: Linear,
}

impl Module for MLP {
    fn forward(&self, input: &Variable) -> Result<Variable> {
        let h = self.fc1.forward(input)?.relu()?;
        self.fc2.forward(&h)
    }

    fn parameters(&self) -> Vec<Parameter> {
        nn::collect_parameters(&[&self.fc1 as &dyn Module, &self.fc2])
    }

    fn name(&self) -> &str { "MLP" }
}

// flodl — free functions, return Variable (differentiable)
let loss = mse_loss(&pred, &target)?;
let loss = cross_entropy_loss(&logits, &labels)?;   // labels: [B] indices or [B,C] one-hot
let loss = nll_loss(&log_probs, &labels)?;           // after log_softmax
let loss = bce_loss(&probs, &target)?;               // from probabilities
let loss = bce_with_logits_loss(&logits, &target)?;  // numerically stable
let loss = l1_loss(&pred, &target)?;
let loss = smooth_l1_loss(&pred, &target, 1.0)?;
let loss = kl_div_loss(&log_probs, &targets)?;
let loss = ctc_loss(&log_probs, &targets, &input_lengths, &target_lengths, 0)?;
let loss = poisson_nll_loss(&pred, &target, true)?;
let loss = focal_loss(&logits, &target, 0.25, 2.0)?;  // alpha, gamma — class imbalance
let loss = triplet_margin_loss(&anchor, &positive, &negative, 1.0)?;
let loss = cosine_embedding_loss(&x1, &x2, &labels, 0.0)?;
let loss = hinge_embedding_loss(&input, &labels, 1.0)?;
let loss = margin_ranking_loss(&x1, &x2, &labels, 0.0)?;

// flodl — optimizers own a clone of the param list
let mut opt = SGD::new(&params, 0.01, 0.9);     // lr, momentum
let mut opt = Adam::new(&params, 0.001);         // lr
let mut opt = AdamW::new(&params, 0.001, 0.01);  // lr, weight_decay
let mut opt = RMSprop::new(&params, 0.01);       // lr (alpha=0.99, eps=1e-8)
let mut opt = Adagrad::new(&params, 0.01);       // lr
let mut opt = RAdam::new(&params, 0.001);        // rectified Adam — auto warmup
let mut opt = NAdam::new(&params, 0.001);        // Nesterov-accelerated Adam

opt.zero_grad();
loss.backward()?;
opt.step()?;  // returns Result<()>

// flodl — builder API
let mut opt = Adam::with_groups()
    .group(&encoder_params, 1e-5)
    .group(&decoder_params, 1e-3)
    .build();

opt.set_group_lr(0, 1e-6);  // adjust one group
opt.set_lr(1e-4);           // adjust all groups

// flodl
for param in &encoder_params {
    param.freeze()?;
}
// Later: param.unfreeze()?;
// Check: param.is_frozen()

// flodl — schedulers produce an lr, you apply it
let sched = StepDecay::new(0.001, 30, 0.1);
let sched = CosineScheduler::new(0.001, 1e-6, 100);
let sched = ExponentialLR::new(0.001, 0.95);
let sched = MultiStepLR::new(0.001, &[30, 60, 90], 0.1);
let sched = OneCycleLR::new(0.01, 1000);                  // 30% warmup
let sched = CyclicLR::new(1e-4, 1e-2, 500);               // symmetric triangle
let lr = sched.lr(step);
opt.set_lr(lr);

// Composable warmup:
let sched = WarmupScheduler::new(CosineScheduler::new(0.001, 1e-6, 100), 0.001, 10);

// Plateau (reduce on metric stall):
let mut sched = PlateauScheduler::new(0.001, 10, 0.1, 1e-6);
let lr = sched.observe(val_loss);
opt.set_lr(lr);

// flodl
let total_norm = clip_grad_norm(&params, 1.0)?;
let max_val = clip_grad_value(&params, 0.5)?;

// flodl — one-call checkpoint (saves params + buffers + structural hash)
model.save_checkpoint("model.fdl")?;
let report = model.load_checkpoint("model.fdl")?;
// report.loaded, report.skipped, report.missing

// Or with any io::Write / io::Read for custom I/O:
let named = model.named_parameters();
let buffers = model.named_buffers();
let hash = Some(model.structural_hash());
save_checkpoint(&mut writer, &named, &buffers, hash)?;
let report = load_checkpoint(&mut reader, &named, &buffers, hash)?;

// Save
model.save_checkpoint("model.fdl")?;
let mut f = File::create("optimizer.fdl")?;
optimizer.save_state(&mut f)?;

// Load
let report = model.load_checkpoint("model.fdl")?;
let mut f = File::open("optimizer.fdl")?;
optimizer.load_state(&mut f)?;

use flodl::nn::{checkpoint_version, migrate_checkpoint_file};

// Check if migration is needed
if checkpoint_version("model.fdl")? < 2 {
    let report = migrate_checkpoint_file(
        "model.fdl",          // old checkpoint (v1)
        "model_v2.fdl",       // migrated output (v2)
        &model.named_parameters(),
        &model.named_buffers(),
    )?;
    println!("{}", report);
    // unchanged (1):
    //   shared/weight
    // remapped (2):
    //   linear_0/weight -> encoder/weight
    //   linear_0/bias -> encoder/bias
}

// Load the migrated checkpoint normally
model.load_checkpoint("model_v2.fdl")?;

// flodl
let device = if cuda_available() { Device::CUDA(0) } else { Device::CPU };

// Move model parameters
module.move_to_device(device);

// Move tensors
let x = x.to_device(device)?;
let x = x.to_device_of(&weights)?;  // match another tensor's device

// Move variables
let x = x.to_device(device)?;

// Create directly on device
let opts = TensorOptions { dtype: DType::Float32, device: Device::CUDA(0) };
let x = Tensor::zeros(&[2, 3], opts)?;

// flodl — returns a new Tensor, then set_data() to apply
let w = xavier_uniform(&[out, inp], inp, out, device)?;
let w = xavier_normal(&[out, inp], inp, out, device)?;
let w = kaiming_uniform(&[out, inp], inp, 0.0, device)?;   // a=0.0 for ReLU
let w = kaiming_normal(&[out, inp], inp, 0.0, device)?;
let w = uniform(&[out, inp], -0.1, 0.1, device)?;
let w = normal(&[out, inp], 0.0, 0.01, device)?;
let w = orthogonal(&[out, inp], 1.0, device)?;             // 2D only
let w = trunc_normal(&[out, inp], 0.0, 0.02, -2.0, 2.0, device)?;  // mean, std, a, b

layer.parameters()[0].set_data(&w);

// flodl
let mut scaler = GradScaler::new();

// Cast model to float16
cast_parameters(&params, DType::Float16);

// Forward + backward with scaled loss
let output = model.forward(&input)?;
let loss = mse_loss(&output, &target)?;
let scaled = scaler.scale(&loss)?;
scaled.backward()?;

// Step with automatic unscaling + inf/nan checking
let stepped = scaler.step(&params, &mut || opt.step())?;
scaler.update();

// flodl
model.train();
for epoch in 0..num_epochs {
    // your data loading here
    opt.zero_grad();
    let output = model.forward(&input)?;
    let loss = cross_entropy_loss(&output, &target)?;
    loss.backward()?;
    clip_grad_norm(&params, 1.0)?;
    opt.step()?;

    opt.set_lr(sched.lr(epoch));
    println!("Epoch {}: loss={:.4}", epoch, loss.item()?);
}

// Inference
model.eval();
let pred = no_grad(|| model.forward(&test_input))?;

fn train_step(model: &Graph, input: &Variable, target: &Variable,
              optimizer: &mut Adam) -> Result<f64> {
    optimizer.zero_grad();
    let output = model.forward(input)?;
    let loss = mse_loss(&output, target)?;
    loss.backward()?;
    optimizer.step()?;
    loss.item()
}

// flodl
let model = FlowBuilder::from(Linear::new(784, 128)?)
    .through(ReLU)
    .through(Linear::new(128, 10)?)
    .build()?;

// flodl — one line
let model = FlowBuilder::from(Linear::new(128, 128)?)
    .also(ReLU)    // skip connection: output = input + ReLU(input)
    .build()?;

// flodl — same block, one builder chain
FlowBuilder::from(prev)
    .also_with(
        downsample_1x1_bn,         // skip branch (Identity if no projection needed)
        conv_bn_relu_conv_bn,      // main branch
    )
    .through(ReLU)
    .build()?;