jaxsne package

Module contents

A jax library for doing dimensionality reduction.

jaxsne is a library for doing dimensionality reduction, with two main methods: scaling for multi-dimensional scaling that tried to match distances directly between spaces, and sne a generic form of t-SNE that instead matches the distribution of distances,

This library uses jax, so users can define their own measures for sne or their own metrics for both, making is robust to different needs. However this library is intended to come with batteries included, so doing scaling(data) or sne(data) will just work.

Remarks

It’s important to note that this implements exact SNE instead of using the optimizations from Barnes-Hut. This means that it is much less performant than the t-SNE provided by scikit-learn, and in general is also less optimized than their exact version, so only use library if you want other metric spaces.

jaxsne.scaling(data: ~jax.Array, n_components: int = 2, *, in_metric: ~collections.abc.Callable[[~jax.Array, ~jax.Array], ~jax.Array] = <PjitFunction of <function euclidean>>, out_metric: ~collections.abc.Callable[[~jax.Array, ~jax.Array], ~jax.Array] = <PjitFunction of <function euclidean>>, max_iter: int = 1024, key: ~jax.Array = Array((), dtype=key<fry>) overlaying: [0 0], gtol: float = 1e-06, xtol: float = 1e-06) → Array

Apply multidimensional scaling.

Multidimensional scaling works by finding points in the reduced dimensional space while trying to keep the pairwise distances as close as possible.

Parameters:

  • data ((n, d)) – The input data, as n rows of d-dimensional points.

  • n_components (The dimension of the space to project into.)

  • in_metric (The metric function for the original data space.)

  • out_metric (The metric function for the destination space.)

  • max_iter (The maximum number of iterations to use to find a solution.)

  • key (A key used for initialization of random projection if the number of) – input points is lower than the input dimension.

  • gtol (Tolerance used for terminating optimization from gradient norm.)

  • xtol (Tolerance used for terminating optimization from trust region size.)

Returns:

result – The reduced dimensional version of data.

Return type:

(n, n_components)

jaxsne.sne(data: ~jax.Array, n_components: int = 2, *, perplexity: float = 30, in_metric: ~collections.abc.Callable[[~jax.Array, ~jax.Array], ~jax.Array] = <PjitFunction of <function euclidean>>, in_measure: ~collections.abc.Callable[[~jax.Array], ~jax.Array] = <PjitFunction of <function gaussian>>, out_metric: ~collections.abc.Callable[[~jax.Array, ~jax.Array], ~jax.Array] = <PjitFunction of <function euclidean>>, out_measure: ~collections.abc.Callable[[~jax.Array], ~jax.Array] = <PjitFunction of <function cauchy>>, out_scale: float | None = None, max_iter: int = 10024, key: ~jax.Array = Array((), dtype=key<fry>) overlaying: [0 0], ptol: float = 1e-06, gtol: float = 1e-06, xtol: float = 1e-06) → Array

Reduce dimension using Stochastic Neighbor Embeddings.

The default implementation of this performs exact tSNE, but alternate metric spaces and distributions can be passed in to perform other variants. Custom metric and distribution functions must be jax jit-able.

Parameters:

  • data ((n, d)) – The input data in “high dimensional” d space.

  • n_components (The dimension of the space to reduce to.)

  • perplexity (This is effectively a smoothing parameter that should be) – adjusted to gain the insights desired. A recommended range is [20, 50]

  • in_metric (The metric space of data.) – (t-SNE = euclidian, vMF-SNE = cosine)

  • in_measure (The distribution to use over the input metric space.) – (t-SNE = gaussian, vMF-SNE = laplace)

  • out_metric (The desired metric space of the output. See information in the) – metrics for you might want to scale these. (t-SNE = euclidian, vMF-SNE = cosine)

  • out_measure (The distribution for matching the output metrics.) – (t-SNE = cauchy, vMF-SNE = laplace)

  • out_scale (It can be beneficial to scale the output metric space to help) – with alignment. This is especially import in more fixed spaces like cosine and poincare, and matters less for euclidian out_metric.

  • max_iter (The maximum number of iterations used to optimize the result.) – This can be adjusted if the optimization is failing for that reason and you’re willing to wait longer.

  • key (The key (seed) used for random projection if necessary to initialize) – the result data. If you’re in this regime and want to try different initializations you can pass key=jax.random.key(seed)

  • ptol (Tolerance used for perplexity matching the input.)

  • gtol (Tolerance used for terminating optimization from gradient norm.)

  • xtol (Tolerance used for terminating optimization from trust region size.)

Returns:

  • result ((n, n_components)) – The resulting data projected into n_components dimensions.

  • Remarks

  • ——-

  • Due to the flexibility of this method, it’s not as readily amenable to the

  • improvements of the Barnes-Hut approximation. As a gneral rule it is also

  • alsower than the scikit-learn’s TSNE in exact mode, so only use this if want

  • alternate metric spaces or distributions.

jaxsne.metric module

Metrics for jaxsne.

A metric function is a distance-like measure between points in a space. SNE only needs a non-negative, symmetric distance that is zero for identical points; it does not rely on the triangle inequality, so metrics need not be proper [metrics](https://en.wikipedia.org/wiki/Metric_space). The shipped cosine, for example, is an angular dissimilarity that violates the triangle inequality. To work with this library a metric must satisfy:

  1. It must be non-negative and symmetric, and zero for identical points (after any projection it applies, see rule 4).

  2. It must treat the last dimension as the dimension of the points.

  3. It must be [jax jit-able](https://docs.jax.dev/en/latest/jit-compilation.html).

  4. It must apply to points in R^d. If it actually applies to some subset of points, it should first project into a space of the same dimension. Such a projection may be undefined at singular points (e.g. cosine at the origin, where the direction is undefined).

jaxsne.metric.cosine(left: Array, right: Array) → Array

Compute the cosine distance.

Before computing the distance, the points are projected onto the l2-ball. It’s normalized to be in [0, 1]

jaxsne.metric.euclidean(left: Array, right: Array) → Array

Compute the euclidian distance.

jaxsne.metric.poincare(left: Array, right: Array) → Array

Compute the hyperbolic distance.

This metric is a two stage process. First, points are projected onto the poincare disk with u = x / (1 + ||x||). Then the standard poincare distance is used on those points.

If data already exists on the poincare-disk, then you need to project it onto the plane with u = x / (1 - ||x||) before using this metric on it.

jaxsne.measure module

Measure for jaxsne.

A measure is a mapping from distances (non-negative) to masses. Since SNE empirically normalizes all distirbutions, these don’t need to be normalized, and should also omit scaling of the distance, even if the standard measure would, as that is also handled by SNE.

  1. They must be proper [measures](https://en.wikipedia.org/wiki/Measure_(mathematics)), which just means they should be non-negative.

  2. They will only be called for positive inputs.

  3. They must be [jax jit-able](https://docs.jax.dev/en/latest/jit-compilation.html).

jaxsne.measure.cauchy(dists: Array) → Array

Unnormalized cauchy distribution.

1 / (1 + x^2)

jaxsne.measure.gaussian(dists: Array) → Array

Unnormalized gaussian distribution.

exp(-x^2)

jaxsne.measure.laplace(dists: Array) → Array

Unnormalized laplace distribution.

exp(-|x|)

Indices and tables