Skip to main content

TensorFlow and t-SNE

TensorFlow 2.21.0 does not provide a maintained high-level t-SNE estimator. A short custom KL loop that uses one Gaussian bandwidth in both spaces omits defining parts of t-SNE. This example keeps TensorFlow at the feature-preparation boundary and delegates the embedding algorithm to scikit-learn.

1. Prepare tensor features​

import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.datasets import load_digits
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE

tf.random.set_seed(42)
digits = load_digits()
X = digits.data[:300] / 16.0
y = digits.target[:300]
X_tensor = tf.convert_to_tensor(X, dtype=tf.float64)
centered = X_tensor - tf.reduce_mean(X_tensor, axis=0, keepdims=True)
assert centered.shape == (300, 64)

2. Fit the maintained embedding​

X_pca = PCA(n_components=30, random_state=42).fit_transform(centered.numpy())
embedding = TSNE(
    n_components=2,
    perplexity=30,
    init="pca",
    learning_rate="auto",
    max_iter=350,
    random_state=42,
).fit_transform(X_pca)
assert embedding.shape == (300, 2)
scatter = plt.scatter(embedding[:, 0], embedding[:, 1], c=y, cmap="tab10", s=16)
plt.colorbar(scatter, label="Digit label")
plt.title("t-SNE view of scikit-learn digits features")
plt.close()

The plotted coordinates are exploratory and transductive. scikit-learn's TSNE has no transform method, so new observations cannot be placed into this fixed map through that estimator.