kooplearn.jax

nn.spectral_contrastive_loss(x, y)	Spectral contrastive loss for self-supervised learning.
nn.vamp_loss(x, y[, schatten_norm, ...])	Variational Approach for learning Markov Processes (VAMP) score.
nn.autoencoder_loss(x, y, x_rec, y_enc, ...)	Single-step Dynamic Autoencoder (DAE) loss.
nn.energy_loss(x, y[, grad_weight])	Energy-based loss function.

NnxFeatureMapEmbedder(encoder[, decoder])

sklearn-style transformer wrapping NNX Modules (encoder/decoder).