Skip to content

api_reference.md

Latest commit

 

History

History
192 lines (143 loc) · 7.74 KB

api_reference.md

File metadata and controls

192 lines (143 loc) · 7.74 KB
Raw

API Reference

import novatorch as nt

Creating tensors

Call Result
nt.tensor(data, dtype=None, device=None, requires_grad=False) from a numpy array, list or scalar; integer input becomes int32
nt.from_numpy(array) alias of tensor
nt.zeros(*shape), nt.ones(*shape) filled tensors
nt.randn(*shape), nt.rand(*shape) normal / uniform
nt.arange(start, stop=None, step=1) range
nt.randint(low, high, shape) int32 tensor
nt.Tensor.full(shape, value) constant
nt.cat(tensors, dim=0), nt.stack(tensors, dim=0) joining

All of them take dtype= and device= and default to float32 on cuda:0.

Tensor

Properties - shape, strides, dtype, device, ndim, numel, requires_grad, grad, grad_fn, is_leaf, name.

Arithmetic - + - * / @ and ** with tensors or scalars, plus the named forms add(other, alpha=1), sub, mul, div, maximum, minimum, pow, sqrt, rsqrt, exp, log, abs, neg, sign, reciprocal, clamp(lo, hi).

Matrix - matmul / @ (2-D, batched, and N-D by folding the leading dimensions), mm, bmm, dot, outer.

Reductions - sum, mean, max, min, argmax, argmin, var, std, norm(p=2, dim=-1, keepdim=False). dim=-1 reduces every dimension; pass an explicit index for a single axis.

Shape - reshape, view, flatten(start, end), squeeze, unsqueeze, transpose(d0, d1), t(), permute(dims), expand(shape), slice(dim, start, end, step=1), narrow(dim, start, length), select(dim, index), index_select(dim, index), contiguous, clone.

Device / dtype - to(device), to(dtype), cuda(idx=0), cpu(), float(), half(), int(), is_cuda(), is_contiguous().

In-place (untracked by autograd; for parameters and buffers) - fill_, zero_, copy_, add_(other, alpha=1), mul_(scalar), normal_(mean, std), uniform_(lo, hi).

Autograd - backward(grad=None, retain_graph=False), detach(), grad (readable and assignable, set to None to clear), zero_grad().

Host access - numpy(), item(), tolist().

Autograd helpers

Call Meaning
nt.no_grad() context manager / decorator that disables graph construction
nt.enable_grad() re-enables it inside a no_grad block
nt.is_grad_enabled() current state
nt.grad(output, inputs, grad_output=None, retain_graph=False) gradients without touching .grad

Layers (nt. or nt.nn.)

Layer Signature
Linear (in_features, out_features, bias=True)
Conv2d (in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)
MaxPool2d / AvgPool2d (kernel_size, stride=kernel_size, padding=0)
AdaptiveAvgPool2d (output_size=1)
LayerNorm (normalized_shape or dim, eps=1e-5, elementwise_affine=True)
BatchNorm2d (num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True)
RMSNorm (dim, eps=1e-6)
Embedding (num_embeddings, embedding_dim, padding_idx=-1, max_norm=0)
Flatten (start_dim=1, end_dim=-1)
Dropout (p=0.5)
Activations ReLU, GELU, SiLU, ELU(alpha), Sigmoid, Tanh, LeakyReLU(slope), Softmax(dim), LogSoftmax(dim)
Sequential (*modules), .add(module), indexable

Module API - parameters(), named_parameters(), state_dict(), load_state_dict(state, strict=True), train(mode=True), eval(), training, zero_grad(set_to_none=True), to(device), cuda(), cpu(), num_parameters().

Subclass nt.nn.Module to write a model in Python; assigned layers and gradient-requiring tensors are registered automatically.

Losses

Loss Target
MSELoss(reduction="mean") float, same shape
L1Loss(reduction="mean") float, same shape
CrossEntropyLoss(label_smoothing=0, ignore_index=-100, reduction="mean") int32 class ids, takes logits
NLLLoss(ignore_index=-100, reduction="mean") int32 class ids, takes log-probabilities
BCEWithLogitsLoss(reduction="mean") float 0/1, takes logits

Optimizers

Optimizer Signature
SGD (params, lr, momentum=0, weight_decay=0, nesterov=False, dampening=0)
Adam (params, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0, amsgrad=False)
AdamW same, weight_decay=0.01, decoupled
RMSprop (params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0)

params accepts a module or a list of tensors. Methods: step(), zero_grad(set_to_none=True), clip_grad_norm(max_norm) (returns the pre-clip norm), get_lr() / set_lr() / .lr, num_steps().

Schedulers (nt.optim.) - StepLR, MultiStepLR, ExponentialLR, CosineAnnealingLR, LinearWarmupCosineDecay, OneCycleLR, ReduceLROnPlateau. All expose step() and get_last_lr().

Functional (nt.functional / a few at nt.)

Activations: relu, sigmoid, tanh, gelu, silu, elu, leaky_relu, softmax(x, dim=-1), log_softmax, dropout(x, p, training).

Layers: linear(x, weight, bias=None), conv2d(x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1), max_pool2d, avg_pool2d, embedding(weight, indices, padding_idx=-1), layer_norm(x, normalized_shape, weight=None, bias=None, eps=1e-5), rms_norm(x, weight=None, eps=1e-6).

Losses: mse_loss, l1_loss, nll_loss, cross_entropy, binary_cross_entropy_with_logits.

Attention: scaled_dot_product_attention(q, k, v, causal=False, dropout_p=0, training=False) (differentiable; picks the fused kernel when no gradient is needed), flash_attention(q, k, v, causal=False) (forward only, head_dim<=64).

Misc: one_hot(indices, num_classes), accuracy(logits, targets), clip_grad_norm(params, max_norm).

Transformer

cfg = nt.GPTConfig()      # vocab_size, max_seq_len, d_model, num_heads,
                          # num_layers, d_ff, dropout, use_flash_attn
model = nt.GPT(cfg)
  • model(tokens) -> logits [B, T, vocab]
  • model.compute_loss(tokens, targets) -> scalar
  • model.generate(tokens, max_new_tokens, temperature=1.0, top_k=0)
  • model.num_params()

Also available: MultiHeadAttention(embed_dim, num_heads, dropout, bias, use_flash) with forward(x, causal), FeedForward, TransformerBlock, TransformerEncoderLayer.

Data

from novatorch.data import DataLoader, TensorDataset, MNISTDataset, CIFAR10Dataset, RandomDataset
loader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False, num_prefetch=2)

Batches come back as GPU tensors; integer arrays stay int32.

Serialization

nt.save(model_or_state, path), nt.load(model, path, strict=True), nt.save_checkpoint(path, model, optimizer=None, epoch=0, extra=None), nt.load_checkpoint(path, model, optimizer=None) (returns the metadata dict). The format is a compressed .npz.

Device and memory

nt.cuda_available(), nt.num_gpus(), nt.set_device(i), nt.current_device(), nt.device("cuda:0"), nt.synchronize(), nt.empty_cache(), nt.device_properties(i), nt.cuda_memory_stats(i), nt.manual_seed(seed).

JIT

nt.jit.execute(expr, **tensors), nt.jit.compile(expr), nt.jit.clear_cache(). Ops: add, sub, mul, div, max, min, relu, gelu, silu, sigmoid, tanh, exp, log, sqrt, abs, neg.

Profiler

nt.profiler.profile(), report(), to_dict(), record_step(loss, lr, throughput), step_records(), gpu_stats(), timer(label), reset().

Limitations

  • Compute is CUDA float32 (int32 for indices). CPU tensors are for I/O.
  • First-order gradients only: no double backward.
  • No mixed-precision autocast, no in-place autograd ops, no sparse tensors.
  • Multi-GPU (distributed/) requires NCCL and is not exercised on Windows.