jaxsne 0.1.0

A library for doing dimensionality reduction with different curvature

jaxsne

A library for doing dimensionality reduction in different metric spaces, or using different distributions.

In addition to regular t-SNE for MNIST

this can also be used to embed points on the sphere

or even into hierarchical hyperbolic space

The downside is that this is generally less performant than the t-SNE provided by scikit-learn, so should only be used if you want to tweak the metrics or measures.

Installation

pip install jaxsne

Basic Usage

import jaxsne

data = ... # n x d
reduced = jaxsne.sne(data)
# or
reduced = jaxsne.scaling(data)

Advanced Usage

import jaxsne
import jax
from jax import Array
from jax import numpy as jnp

@jax.jit
def manhattan(left: Array, right: Array) -> Array:
    return jnp.abs(left - right).sum(-1)


data = ... # n x d
reduced = jaxsne.sne(data, in_metric=manhattan, out_metric=manhattan)

Development

uv run ruff format
uv run ruff check
uv run mypy jaxsne
uv run pytest

Tasks

• Extend to Barnes-Hut-SNE