Casper/factored processes

simplexity/activations/activation_analyses.py

@@ -35,7 +35,7 @@ def analyze(
         self,
         activations: Mapping[str, jax.Array],
         weights: jax.Array,
-        belief_states: jax.Array | None = None,
+        belief_states: jax.Array | tuple[jax.Array, ...] | None = None,
     ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]:
         """Analyze activations and return scalar metrics and projections."""
         ...
@@ -76,13 +76,14 @@ def __init__(
         concat_layers: bool = False,
         use_probs_as_weights: bool = True,
         fit_intercept: bool = True,
+        to_factors: bool = False,
     ) -> None:
         super().__init__(
             analysis_type="linear_regression",
             last_token_only=last_token_only,
             concat_layers=concat_layers,
             use_probs_as_weights=use_probs_as_weights,
-            analysis_kwargs={"fit_intercept": fit_intercept},
+            analysis_kwargs={"fit_intercept": fit_intercept, "to_factors": to_factors},
         )
 
 
@@ -97,8 +98,9 @@ def __init__(
         use_probs_as_weights: bool = True,
         rcond_values: Sequence[float] | None = None,
         fit_intercept: bool = True,
+        to_factors: bool = False,
     ) -> None:
-        analysis_kwargs: dict[str, Any] = {"fit_intercept": fit_intercept}
+        analysis_kwargs: dict[str, Any] = {"fit_intercept": fit_intercept, "to_factors": to_factors}
         if rcond_values is not None:
             analysis_kwargs["rcond_values"] = tuple(rcond_values)
         super().__init__(

simplexity/activations/activation_tracker.py

@@ -20,7 +20,7 @@ class PreparedActivations:
     """Prepared activations with belief states and sample weights."""
 
     activations: Mapping[str, jax.Array]
-    belief_states: jax.Array | None
+    belief_states: jax.Array | tuple[jax.Array, ...] | None
     weights: jax.Array
 
 
@@ -48,16 +48,26 @@ def _to_jax_array(value: Any) -> jax.Array:
     return jnp.asarray(value)
 
 
+def _convert_tuple_to_jax_array(value: tuple[Any, ...]) -> tuple[jax.Array, ...]:
+    """Convert a tuple of supported tensor types to JAX arrays."""
+    return tuple(_to_jax_array(v) for v in value)
+
+
 def prepare_activations(
     inputs: jax.Array | torch.Tensor | np.ndarray,
-    beliefs: jax.Array | torch.Tensor | np.ndarray,
+    beliefs: jax.Array
+    | torch.Tensor
+    | np.ndarray
+    | tuple[jax.Array, ...]
+    | tuple[torch.Tensor, ...]
+    | tuple[np.ndarray, ...],
     probs: jax.Array | torch.Tensor | np.ndarray,
     activations: Mapping[str, jax.Array | torch.Tensor | np.ndarray],
     prepare_options: PrepareOptions,
 ) -> PreparedActivations:
     """Preprocess activations by deduplicating sequences, selecting tokens/layers, and computing weights."""
     inputs = _to_jax_array(inputs)
-    beliefs = _to_jax_array(beliefs)
+    beliefs = _convert_tuple_to_jax_array(beliefs) if isinstance(beliefs, tuple) else _to_jax_array(beliefs)
     probs = _to_jax_array(probs)
     activations = {name: _to_jax_array(layer) for name, layer in activations.items()}
 
@@ -74,7 +84,7 @@ def prepare_activations(
     weights = (
         dataset.probs
         if prepare_options.use_probs_as_weights
-        else _get_uniform_weights(belief_states.shape[0], belief_states.dtype)
+        else _get_uniform_weights(dataset.probs.shape[0], dataset.probs.dtype)
     )
 
     if prepare_options.concat_layers:
@@ -98,7 +108,12 @@ def __init__(self, analyses: Mapping[str, ActivationAnalysis]):
     def analyze(
         self,
         inputs: jax.Array | torch.Tensor | np.ndarray,
-        beliefs: jax.Array | torch.Tensor | np.ndarray,
+        beliefs: jax.Array
+        | torch.Tensor
+        | np.ndarray
+        | tuple[jax.Array, ...]
+        | tuple[torch.Tensor, ...]
+        | tuple[np.ndarray, ...],
         probs: jax.Array | torch.Tensor | np.ndarray,
         activations: Mapping[str, jax.Array | torch.Tensor | np.ndarray],
     ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]:

simplexity/analysis/layerwise_analysis.py

@@ -34,21 +34,23 @@ class AnalysisRegistration:
 
 def _validate_linear_regression_kwargs(kwargs: Mapping[str, Any] | None) -> dict[str, Any]:
     provided = dict(kwargs or {})
-    allowed = {"fit_intercept"}
+    allowed = {"fit_intercept", "to_factors"}
     unexpected = set(provided) - allowed
     if unexpected:
         raise ValueError(f"Unexpected linear_regression kwargs: {sorted(unexpected)}")
     fit_intercept = bool(provided.get("fit_intercept", True))
-    return {"fit_intercept": fit_intercept}
+    to_factors = bool(provided.get("to_factors", False))
+    return {"fit_intercept": fit_intercept, "to_factors": to_factors}
 
 
 def _validate_linear_regression_svd_kwargs(kwargs: Mapping[str, Any] | None) -> dict[str, Any]:
     provided = dict(kwargs or {})
-    allowed = {"fit_intercept", "rcond_values"}
+    allowed = {"fit_intercept", "rcond_values", "to_factors"}
     unexpected = set(provided) - allowed
     if unexpected:
         raise ValueError(f"Unexpected linear_regression_svd kwargs: {sorted(unexpected)}")
     fit_intercept = bool(provided.get("fit_intercept", True))
+    to_factors = bool(provided.get("to_factors", False))
     rcond_values = provided.get("rcond_values")
     if rcond_values is not None:
         if not isinstance(rcond_values, (list, tuple)):
@@ -58,6 +60,7 @@ def _validate_linear_regression_svd_kwargs(kwargs: Mapping[str, Any] | None) ->
         rcond_values = tuple(float(v) for v in rcond_values)
     return {
         "fit_intercept": fit_intercept,
+        "to_factors": to_factors,
         "rcond_values": rcond_values,
     }
 
@@ -152,7 +155,7 @@ def analyze(
         self,
         activations: Mapping[str, jax.Array],
         weights: jax.Array,
-        belief_states: jax.Array | None = None,
+        belief_states: jax.Array | tuple[jax.Array, ...] | None = None,
     ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]:
         """Analyze activations and return namespaced scalar metrics and projections."""
         if self._requires_belief_states and belief_states is None:

simplexity/analysis/linear_regression.py

@@ -139,22 +139,56 @@ def linear_regression_svd(
 def layer_linear_regression(
     layer_activations: jax.Array,
     weights: jax.Array,
-    belief_states: jax.Array | None,
+    belief_states: jax.Array | tuple[jax.Array, ...] | None,
+    to_factors: bool = False,
     **kwargs: Any,
 ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]:
     """Layer-wise regression helper that wraps :func:`linear_regression`."""
     if belief_states is None:
         raise ValueError("linear_regression requires belief_states")
-    return linear_regression(layer_activations, belief_states, weights, **kwargs)
+
+    if to_factors:
+        scalars, projections = {}, {}
+        if not isinstance(belief_states, tuple):
+            raise ValueError("belief_states must be a tuple when to_factors is True")
+        for factor_idx, factor in enumerate(belief_states):
+            if not isinstance(factor, jax.Array):
+                raise ValueError("Each factor in belief_states must be a jax.Array")
+            factor_scalars, factor_projections = linear_regression(layer_activations, factor, weights, **kwargs)
+            for key, value in factor_scalars.items():
+                scalars[f"factor_{factor_idx}/{key}"] = value
+            for key, value in factor_projections.items():
+                projections[f"factor_{factor_idx}/{key}"] = value
+        return scalars, projections
+    else:
+        belief_states = jnp.concatenate(belief_states, axis=-1) if isinstance(belief_states, tuple) else belief_states
+        return linear_regression(layer_activations, belief_states, weights, **kwargs)
 
 
 def layer_linear_regression_svd(
     layer_activations: jax.Array,
     weights: jax.Array,
-    belief_states: jax.Array | None,
+    belief_states: jax.Array | tuple[jax.Array, ...] | None,
+    to_factors: bool = False,
     **kwargs: Any,
 ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]:
     """Layer-wise regression helper that wraps :func:`linear_regression_svd`."""
     if belief_states is None:
         raise ValueError("linear_regression_svd requires belief_states")
-    return linear_regression_svd(layer_activations, belief_states, weights, **kwargs)
+
+    if to_factors:
+        scalars, projections = {}, {}
+        if not isinstance(belief_states, tuple):
+            raise ValueError("belief_states must be a tuple when to_factors is True")
+        for factor_idx, factor in enumerate(belief_states):
+            if not isinstance(factor, jax.Array):
+                raise ValueError("Each factor in belief_states must be a jax.Array")
+            factor_scalars, factor_projections = linear_regression_svd(layer_activations, factor, weights, **kwargs)
+            for key, value in factor_scalars.items():
+                scalars[f"factor_{factor_idx}/{key}"] = value
+            for key, value in factor_projections.items():
+                projections[f"factor_{factor_idx}/{key}"] = value
+        return scalars, projections
+    else:
+        belief_states = jnp.concatenate(belief_states, axis=-1) if isinstance(belief_states, tuple) else belief_states
+        return linear_regression_svd(layer_activations, belief_states, weights, **kwargs)