xax 0.3.7__py3-none-any.whl → 0.3.8__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.7"
+__version__ = "0.3.8"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -31,6 +31,10 @@ __all__ = [
     "TransformerBlock",
     "TransformerCache",
     "TransformerStack",
+    "Categorical",
+    "Distribution",
+    "MixtureOfGaussians",
+    "Normal",
     "FourierEmbeddings",
     "IdentityPositionalEmbeddings",
     "LearnedPositionalEmbeddings",
@@ -219,6 +223,10 @@ NAME_MAP: dict[str, str] = {
     "TransformerBlock": "nn.attention",
     "TransformerCache": "nn.attention",
     "TransformerStack": "nn.attention",
+    "Categorical": "nn.distributions",
+    "Distribution": "nn.distributions",
+    "MixtureOfGaussians": "nn.distributions",
+    "Normal": "nn.distributions",
     "FourierEmbeddings": "nn.embeddings",
     "IdentityPositionalEmbeddings": "nn.embeddings",
     "LearnedPositionalEmbeddings": "nn.embeddings",
@@ -405,6 +413,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         TransformerCache,
         TransformerStack,
     )
+    from xax.nn.distributions import Categorical, Distribution, MixtureOfGaussians, Normal
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,

xax/nn/distributions.py

@@ -0,0 +1,181 @@
+"""Defines some probability distribution helper functions.
+
+In general, it is preferrable to use Distrax or another library, but we wanted
+to have a simple interface of our own so that we can quickly upgrade Jax
+versions (since Distrax is tied pretty closely to Tensorflow).
+"""
+
+__all__ = [
+    "Distribution",
+    "Categorical",
+    "Normal",
+    "MixtureOfGaussians",
+]
+
+from abc import ABC, abstractmethod
+
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, PRNGKeyArray
+
+
+class Distribution(ABC):
+    @abstractmethod
+    def log_prob(self, x: Array) -> Array: ...
+
+    @abstractmethod
+    def sample(self, key: PRNGKeyArray) -> Array: ...
+
+    @abstractmethod
+    def mode(self) -> Array: ...
+
+    @abstractmethod
+    def entropy(self) -> Array: ...
+
+
+class Categorical(Distribution):
+    def __init__(self, logits_n: Array) -> None:
+        self.logits_n = logits_n
+
+    @property
+    def num_categories(self) -> int:
+        return self.logits_n.shape[-1]
+
+    def log_prob(self, x: Array) -> Array:
+        """Compute log probability for specific categories.
+
+        Args:
+            x: Array of category indices
+
+        Returns:
+            Log probabilities for the given categories
+        """
+        log_probs = jax.nn.log_softmax(self.logits_n, axis=-1)
+        # Use advanced indexing to get the log probabilities for the given categories
+        return log_probs[x]
+
+    def sample(self, key: PRNGKeyArray) -> Array:
+        return jax.random.categorical(key, self.logits_n, axis=-1)
+
+    def mode(self) -> Array:
+        return self.logits_n.argmax(axis=-1)
+
+    def entropy(self) -> Array:
+        """Compute entropy of the categorical distribution."""
+        probs = jax.nn.softmax(self.logits_n, axis=-1)
+        log_probs = jax.nn.log_softmax(self.logits_n, axis=-1)
+        return -jnp.sum(probs * log_probs, axis=-1)
+
+
+class Normal(Distribution):
+    def __init__(self, loc: Array, scale: Array) -> None:
+        self.loc = loc
+        self.scale = scale
+
+    def log_prob(self, x: Array) -> Array:
+        return -0.5 * jnp.log(2 * jnp.pi) - jnp.log(self.scale) - (x - self.loc) ** 2 / (2 * self.scale**2)
+
+    def sample(self, key: PRNGKeyArray) -> Array:
+        return self.loc + self.scale * jax.random.normal(key, self.loc.shape)
+
+    def mode(self) -> Array:
+        return self.loc
+
+    def entropy(self) -> Array:
+        return jnp.log(2 * jnp.pi * jnp.e) + jnp.log(self.scale)
+
+
+class MixtureOfGaussians(Distribution):
+    def __init__(self, means_nm: Array, stds_nm: Array, logits_nm: Array) -> None:
+        """Initialize a mixture of Gaussians.
+
+        Args:
+            means_nm: Array of shape (..., n_components) containing means
+            stds_nm: Array of shape (..., n_components) containing standard deviations
+            logits_nm: Array of shape (..., n_components) containing mixing logits
+        """
+        self.means_nm = means_nm
+        self.stds_nm = stds_nm
+        self.logits_nm = logits_nm
+
+    def log_prob(self, x: Array) -> Array:
+        """Compute log probability of the mixture.
+
+        Args:
+            x: Array of shape (...,) containing values to evaluate
+
+        Returns:
+            Log probabilities of shape (...,)
+        """
+        # Expand x to match component dimensions
+        x_expanded = x[..., None]  # Shape: (..., 1)
+
+        # Compute log probabilities for each component
+        component_log_probs = (
+            -0.5 * jnp.log(2 * jnp.pi)
+            - jnp.log(self.stds_nm)
+            - (x_expanded - self.means_nm) ** 2 / (2 * self.stds_nm**2)
+        )
+
+        # Compute mixing weights
+        mixing_logits = jax.nn.log_softmax(self.logits_nm, axis=-1)
+
+        # Combine using log-sum-exp trick for numerical stability
+        return jax.scipy.special.logsumexp(component_log_probs + mixing_logits, axis=-1)
+
+    def sample(self, key: PRNGKeyArray) -> Array:
+        """Sample from the mixture of Gaussians.
+
+        Args:
+            key: PRNG key
+
+        Returns:
+            Samples of shape (...,) where ... are the batch dimensions
+        """
+        # Sample component indices
+        component_key, sample_key = jax.random.split(key)
+        component_indices = jax.random.categorical(component_key, self.logits_nm, axis=-1)
+
+        # Sample from selected components using advanced indexing
+        # We need to handle the case where we have batch dimensions
+        batch_shape = self.means_nm.shape[:-1]  # All dimensions except the last (components)
+
+        # Reshape for easier indexing
+        means_flat = self.means_nm.reshape(-1, self.means_nm.shape[-1])
+        stds_flat = self.stds_nm.reshape(-1, self.stds_nm.shape[-1])
+        indices_flat = component_indices.reshape(-1)
+
+        # Get selected means and stds
+        selected_means = means_flat[jnp.arange(len(indices_flat)), indices_flat]
+        selected_stds = stds_flat[jnp.arange(len(indices_flat)), indices_flat]
+
+        # Generate random noise
+        noise = jax.random.normal(sample_key, selected_means.shape)
+
+        # Reshape back to original batch shape
+        samples = selected_means + selected_stds * noise
+        return samples.reshape(batch_shape)
+
+    def mode(self) -> Array:
+        """Return the mode of the mixture (approximate - returns mean of highest weight component)."""
+        mixing_weights = jax.nn.softmax(self.logits_nm, axis=-1)
+        max_weight_idx = jnp.argmax(mixing_weights, axis=-1)
+
+        # Use advanced indexing to get the means of the highest weight components
+        batch_shape = self.means_nm.shape[:-1]
+        means_flat = self.means_nm.reshape(-1, self.means_nm.shape[-1])
+        indices_flat = max_weight_idx.reshape(-1)
+
+        selected_means = means_flat[jnp.arange(len(indices_flat)), indices_flat]
+        return selected_means.reshape(batch_shape)
+
+    def entropy(self) -> Array:
+        """Compute entropy of the mixture (approximate)."""
+        mixing_weights = jax.nn.softmax(self.logits_nm, axis=-1)
+        component_entropies = jnp.log(2 * jnp.pi * jnp.e) + jnp.log(self.stds_nm)
+
+        # Weighted sum of component entropies plus mixing entropy
+        weighted_entropies = jnp.sum(mixing_weights * component_entropies, axis=-1)
+        mixing_entropy = -jnp.sum(mixing_weights * jnp.log(mixing_weights + 1e-8), axis=-1)
+
+        return weighted_entropies + mixing_entropy

{xax-0.3.7.dist-info → xax-0.3.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.3.7
+Version: 0.3.8
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.3.7.dist-info → xax-0.3.8.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=YCDjLRwliJCyEmNFC56PNQXV9Vn9Fr13VJS_am4h3To,16336
+xax/__init__.py,sha256=5NGaVm9X36LhG-Tl1hc7Lk1SmnTZvyu8G1iFDixpqLc,16665
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
@@ -9,6 +9,7 @@ xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
 xax/core/state.py,sha256=_gtINsRc310Bu_HuIYsDoOKTZa6DgU2tz0IOKkdnY9Q,3813
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/attention.py,sha256=m6yEoRqf7-wLgrEltaR6CxF_Cody0MaNtAkuKk39qJI,31176
+xax/nn/distributions.py,sha256=096IDvoJ0ZA4SqcfgNSmrICsGcsKVcTAh0Vl6SwN3-o,6343
 xax/nn/embeddings.py,sha256=8tAuAPdkVj-U5IwtRZKHA0WYMFRbpCuwyAxcChdKhbE,11784
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
 xax/nn/geom.py,sha256=c9K52vLm-V-15CRqMNx0OmqsWfb3PHQxXW4OSx9kCAk,10635
@@ -59,9 +60,9 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.3.7.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.3.7.dist-info/METADATA,sha256=8Zb0pvTJOjrCHK7giM2MbhlGCPREQewJK3GgRDQNWY0,1246
-xax-0.3.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.3.7.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.3.7.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.3.7.dist-info/RECORD,,
+xax-0.3.8.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.8.dist-info/METADATA,sha256=d4UVJYHBKGAJTdC8G4IHt9kI44lbexOWIiZnkICd0pM,1246
+xax-0.3.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.8.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.8.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.8.dist-info/RECORD,,