feat: add reusable transformer, checkpoint, and SafeTensors support

[antlobach]

Changes

Generic framework functionality extracted while porting TabICLv1 to Clojure. The TabICL model itself lives in antlobach/tabicl-clorch, which consumes these APIs.

Reusable transformer modules (clorch.nn)

• multihead-attention — batch-first multi-head attention with packed QKV projection and PyTorch-equivalent parameter layout. Forward accepts a tensor for self-attention or a map with :query/:key/:value, :attention-mask, :key-padding-mask, :causal?, :rope, :cached-key/:cached-value, and :return-kv?.
• transformer-encoder-layer — pre-norm (:norm-first? true) or post-norm block with :relu/:gelu/custom activation, matching torch.nn.TransformerEncoderLayer semantics.
• transformer-encoder — layer stack mirroring torch.nn.TransformerEncoder, optionally returning a per-layer key/value cache for incremental decoding.

Model lifecycle fixes required by GPU execution

• defmodel now uses keyword field keys when recursively applying train and to; the previous generated methods looked fields up by their values and therefore did nothing.
• defmodel device/dtype conversion now delegates every field through nn/to, including direct Parameter fields.
• Moving a custom Parameter now detaches the converted tensor and restores requires-grad, preserving leaf status so optimizers receive gradients after CPU/CUDA or dtype conversion.

Functional and tensor support

• clorch.nn.functional/combine-attention-masks — composes SDPA masks with module-convention key-padding masks through boolean and additive-float paths.
• torch/precompute-rope-freqs and torch/apply-rope support :interleaved? for adjacent-pair rotation, used by Llama-family and TabICL models; the default half-split layout is unchanged.
• torch/precompute-rope-freqs accepts :device, allowing RoPE tables to be created for CUDA inputs without a device mismatch.

SafeTensors I/O (clorch.safetensors)

• load-file, save-file!, metadata, and load-state-dict! read/write SafeTensors archives without Python or pickle execution.
• Supports all SafeTensors specification dtypes, validates headers and offsets, writes atomically, and loads flat dotted parameter names into defmodel models.

Checkpoint and initialization primitives

• nn.init/trunc-normal! — inverse-CDF truncated-normal initialization.
• nn/load-flat-state-dict! / nn/flat-parameter-names — dotted-path state-dict loading for native modules, records, parameters, tensors, and nested collections.
• Primitive float[] tensor construction for binary checkpoint readers.
• Fixed torch/copy-: it used set_ (storage rebinding) rather than value copying.

Verification

Check	Result
Focused transformer/SafeTensors/model tests on Linux GPU host, JVM 25	10 tests / 40 assertions, 0 failures
Full GPU release suite (scripts/release-check.sh --mode gpu)	175 tests / 773 assertions, 0 failures; reflection check clean
CUDA/NCCL/DDP/AMP/fused-attention/checkpoint smoke	Passed
Explicit CUDA multihead-attention forward/backward	CUDA output and input gradient; finite values
Explicit CUDA transformer-encoder-layer forward/backward	CUDA output; all 12 parameters receive finite CUDA gradients
Explicit CUDA interleaved RoPE	CUDA output [2 4 2 8]; finite values
SafeTensors CUDA round trip	Both tensors restored to CUDA; max error 0.0
TabICL consumer GPU smoke	CUDA logits [2 3 3]; 94 parameters; all finite
Numerical parity vs PyTorch 2.1.2 (run_transformer_comparison.sh)	self-attention 2.98e-8, cross-attention 5.96e-8, transformer layer 1.34e-6, interleaved RoPE 0.0 (tolerance 2e-5)
CPU benchmark vs PyTorch 2.1.2 (run_transformer_benchmark.sh, dim 512 / 8 heads / batch 4 / seq 256 / 1 thread)	Clorch 32.0 ms vs PyTorch 35.7 ms median
Local focused regression suite	18 tests / 54 assertions, 0 failures
clj-kondo on changed Clojure files	0 errors, 0 warnings
cljfmt	clean
MkDocs strict build	passes

GPU environment: Ubuntu 24.04, OpenJDK 25.0.4.1, NVIDIA RTX 2000 Ada 16 GB, driver 580.159.04, CUDA 13.1 user-space libraries, LibTorch 2.10.0-1.5.13.

Consumer verification is in antlobach/tabicl-clorch#1. It replaces the port's private attention, transformer, RoPE, and SafeTensors implementations with these APIs.

Files

• Source: src/clorch/{nn.clj,nn/functional.clj,nn/init.clj,torch.clj,safetensors.clj}
• Tests: test/clorch/{nn_transformer_test.clj,safetensors_test.clj,nn_init_test.clj,nn_model_test.clj,torch_test.clj}
• Numerical comparison: tests_comparison/{compare_transformer.clj,transformer_reference.py}, run_transformer_comparison.sh
• Benchmark: benchmarks/transformer_{clorch.clj,pytorch.py}, run_transformer_benchmark.sh
• Documentation: README.md, PORTING_STATUS.md, docs/{nn,functional,init,tensors,safetensors,pytorch-parity}.md, docs/mkdocs.yml