From f25015b1ecc3d7cda2bc90a547433457d92dd26d Mon Sep 17 00:00:00 2001
From: Faezeh Habibi <155960330+Faezehabibi@users.noreply.github.com>
Date: Wed, 27 May 2026 16:55:58 -0400
Subject: [PATCH] Implement effective rank calculation in effective_dim.py

Adds a function to calculate the effective rank of a code matrix using Shannon entropy.
---
 ngclearn/utils/analysis/effective_dim.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/ngclearn/utils/analysis/effective_dim.py b/ngclearn/utils/analysis/effective_dim.py
index a7b5aca0..7b46f318 100644
--- a/ngclearn/utils/analysis/effective_dim.py
+++ b/ngclearn/utils/analysis/effective_dim.py
@@ -20,3 +20,27 @@ def participation_ratio(latent_codes):
 
     return tr2_cov / cov2_tr if cov2_tr > 0 else float("nan")
 
+
+
+
+def rankme(Z, eps=1e-7):
+    """
+    Calculates the effective rank of for a code matrix Z
+    effective rank = exp(Shannon entropy), from Garrido, Balestriero,
+    Najman & LeCun, "RankMe: Assessing the Downstream Performance of Pretrained
+    Self-Supervised Representations by Their Rank" (ICML 2023, arXiv:2210.02885).
+
+    Args:
+        latent_codes: a set of (N x D) latent code vectors (one row per vector code)
+
+    Returns:
+        scalar measurement of the effective dimension
+    """
+
+    singular_values = jnp.linalg.svd(Z, compute_uv=False)   ## singular values of Z
+    sum_singular_vals = jnp.sum(singular_values)            ## L1
+    if sum_singular_vals <= 0:
+        return float("nan")
+    p = singular_values / sum_singular_vals + eps          ## L1-normalized singular value
+    shannon_entropy = -jnp.sum(p * jnp.log(p))             ## Shannon entropy
+    return jnp.exp(shannon_entropy)                        ## exp(Shannon entropy) = effective rank