Crate rlx_umap 

Parametric UMAP on RLX — k-NN building blocks + full Umap::fit API matching fast-umap.

• Umap::fit / fit_with_progress — sparse cross-entropy training
• FittedUmap::transform — inference with training z-score stats
• FittedUmap::save / [load] — safetensors or GGUF (.ruama load only for legacy files)

Call register once per process before Session execution.

Quick start

rlx_umap::register();

use rlx_ir::{DType, Graph, Shape};
use rlx_umap::knn_indices_and_distances;

let mut g = Graph::new("knn");
let pairwise = g.input("pairwise", Shape::new(&[64, 64], DType::F32));
let (idx, dist) = knn_indices_and_distances(&mut g, pairwise, 5);
g.set_outputs(vec![idx, dist]);

Re-exports

pub use knn_attrs::KnnAttrs;

pub use pack::unpack_knn_packed;

pub use pairwise::cosine_pairwise_reference;

pub use pairwise::euclidean_pairwise_reference;

pub use parity::KnnParityReport;

pub use parity::compare_knn;

pub use parity::max_pairwise_error;

pub use ops::UMAP_KNN;

pub use ops::UMAP_KNN_BWD;

pub use ops::register_umap_ops;

pub use graph::cosine_knn_graph;

pub use graph::cosine_knn_packed_graph;

pub use graph::knn_graph;

pub use graph::knn_indices_and_distances;

pub use graph::pairwise_cosine_graph;

pub use graph::pairwise_euclidean_graph;

pub use graph::split_knn_packed;

pub use config::UmapConfig;

pub use data::load_csv;

pub use data::load_f64_matrix;

pub use data::load_synthetic;

pub use data::write_embedding_csv;

pub use fitted::FittedUmap;

pub use model_io::EXT_GGUF;

pub use model_io::EXT_RUAMA;

pub use model_io::EXT_SAFETENSORS;

pub use model_io::MODEL_EXT;

pub use model_io::format_from_path;

pub use model_io::model_path;

pub use model_io::model_path_with_ext;

pub use model_io::weight_shapes;

pub use serialize::LoadedModel;

pub use serialize::ModelMetadata;

pub use serialize::SaveBundle;

pub use serialize::load_model;

pub use serialize::load_weights;

pub use serialize::save_model;

pub use serialize::save_weights;

pub use train::EpochProgress;

pub use training::FitOptions;

pub use training::TrainResult;

pub use training::fit;

pub use training::fit_with_progress;

pub use training::train_only;

pub use umap::Umap;

pub use utils::NormStats;

pub use weights::WeightStore;

Modules

adam

Host-side Adam optimizer (matches fast-umap defaults).

config

UMAP configuration (compatible with fast-umap / umap-rs style).

data

Load training data for the train-umap binary and examples.

encoder

fitted

Fitted parametric UMAP model.

graph

Graph builders for UMAP pairwise distances and k-NN.

interrupt

Ctrl-C graceful cancellation (fast-umap compatible).

knn

Reference k-NN — implemented in rlx_cpu::umap_knn for sharing with GPU hosts.

knn_attrs

Attribute blob for umap.knn / umap.knn_backward (k neighbours per row).

model

Compiled UMAP model graphs.

model_io

Model I/O in safetensors (default) and GGUF (F32), with legacy .ruama load.

nn_descent

NN-Descent approximate k-NN (ported from fast-umap).

ops

RLX custom-op registration for UMAP k-NN.

pack

Packed [n, 2k] layout helpers (umap.knn output).

pairwise

Reference pairwise distance matrices (row-major [n, n]).

parity

Parity metrics vs a reference k-NN run.

pca

PCA for warm-starting parametric UMAP embeddings (host-only, no Burn).

prelude

Common re-exports for parametric UMAP on RLX.

serialize

Weight and model I/O.

session

Session helpers for backends with custom-op limitations.

train

Sparse UMAP training loop (port of fast-umap train_sparse).

training

Parametric UMAP training on RLX Session + autodiff (no Burn).

umap

High-level Umap::fit API (parametric UMAP on RLX).

utils

Data normalization and layout helpers.

weights

Named parameter tensors for compiled UMAP graphs.

Enums

Device

Target device for graph execution.

Functions

knn_backward_pairwise

Backward: scatter d_dist [n, k] into d_pairwise [n, n].

knn_forward_packed

Forward k-NN: for each row, select k smallest off-diagonal distances.

register

Register UMAP custom ops (IR + CPU). Alias of register_umap_ops.