vsax 1.2.0

JAX-native Vector Symbolic Algebra library

VSAX: Vector Symbolic Algebra for JAX

VSAX is a GPU-accelerated, JAX-native Python library for Vector Symbolic Architectures (VSAs). It provides composable symbolic representations using hypervectors, algebraic operations for binding and bundling, and encoding strategies for symbolic and structured data.

Features

• 🚀 Three VSA Models: FHRR, MAP, and Binary implementations ✅
• 🏭 Factory Functions: One-line model creation with sensible defaults ✅
• 💾 VSAMemory: Dictionary-style symbol management ✅
• 📊 5 Core Encoders: Scalar, Sequence, Set, Dict, and Graph encoders ✅
• 🎨 Custom Encoders: Easy-to-extend AbstractEncoder base class ✅
• 🔍 Similarity Metrics: Cosine, dot, and Hamming similarity ✅
• ⚡ Batch Operations: GPU-accelerated vmap operations for parallel processing ✅
• 💾 I/O & Persistence: Save/load basis vectors to JSON ✅
• 🎮 GPU Utilities: Device management, benchmarking, CPU/GPU comparison ✅
• 🔧 Clifford Operators: Exact, compositional, invertible transformations for reasoning ✅ NEW in v1.1.0
• 📍 Fractional Power Encoding (FPE): Continuous value encoding using v^r for spatial/function representations ✅ NEW in v1.2.0
• 🗺️ Spatial Semantic Pointers (SSP): Continuous spatial encoding with what/where queries ✅ NEW in v1.2.0
• 🎯 Vector Function Architecture (VFA): Function approximation in RKHS for density estimation, regression, image processing ✅ NEW in v1.2.0
• 🚀 GPU-Accelerated: Built on JAX for automatic GPU acceleration (5-30x speedup)
• 🧩 Modular Architecture: Clean separation between representations and operations
• 🧬 Complete Representations: Complex, Real, and Binary hypervectors ✅
• ⚙️ Full Operation Sets: FFT-based FHRR, MAP, and XOR/majority Binary ops ✅
• 🎲 Random Sampling: Sampling utilities for all representation types ✅
• 📚 Comprehensive Documentation: Full API docs and examples ✅
• 📓 Interactive Tutorials: Jupyter notebooks with real datasets (MNIST, knowledge graphs) ✅ NEW in v0.7.1
• ✅ 95% Test Coverage: 450 tests ensuring reliability

Installation

From PyPI (Recommended)

pip install vsax

Or with uv:

uv pip install vsax

From Source

Using uv (Recommended)

uv is a fast Python package installer and resolver. Install it first:

# Install uv (Unix/macOS)
curl -LsSf https://astral.sh/uv/install.sh | sh

# Install uv (Windows)
powershell -c "irm https://astral.sh/uv/install.ps1 | iex"

Then install VSAX:

git clone https://github.com/vasanthsarathy/vsax.git
cd vsax

# Create virtual environment and install package
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
uv pip install -e .

Using pip

git clone https://github.com/vasanthsarathy/vsax.git
cd vsax
pip install -e .

Development Installation

Using uv (Recommended)

uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
uv pip install -e ".[dev,docs]"

Using pip

pip install -e ".[dev,docs]"

Quick Start

New in v0.6.0: Save and load basis vectors!

Simple Example

from vsax import create_fhrr_model, VSAMemory, DictEncoder, save_basis, load_basis
from vsax.similarity import cosine_similarity
from vsax.utils import vmap_bind

# Create model with factory function (one line!)
model = create_fhrr_model(dim=512)

# Create memory for symbol management
memory = VSAMemory(model)

# Add symbols - automatically samples and stores hypervectors
memory.add_many(["subject", "action", "dog", "run", "cat", "jump"])

# Dictionary-style access
dog = memory["dog"]

# Encode structured data with DictEncoder
encoder = DictEncoder(model, memory)
sentence = encoder.encode({"subject": "dog", "action": "run"})

# Bind concepts (circular convolution)
dog_is_animal = model.opset.bind(dog.vec, memory["animal"].vec)

# Bundle concepts (sum and normalize)
pets = model.opset.bundle(memory["dog"].vec, memory["cat"].vec)

# NEW: Similarity search
similarity = cosine_similarity(memory["dog"], memory["cat"])
print(f"Dog-Cat similarity: {similarity:.3f}")

# NEW: Batch operations (GPU-accelerated)
import jax.numpy as jnp
nouns = jnp.stack([memory["dog"].vec, memory["cat"].vec])
verbs = jnp.stack([memory["run"].vec, memory["jump"].vec])
actions = vmap_bind(model.opset, nouns, verbs)  # Parallel binding!

# NEW: Save and load basis vectors
save_basis(memory, "my_basis.json")  # Persist to JSON
memory_new = VSAMemory(model)
load_basis(memory_new, "my_basis.json")  # Load from JSON

All Three Models

VSAX supports three VSA models, all with the same simple API:

from vsax import create_fhrr_model, create_map_model, create_binary_model, VSAMemory

# FHRR: Complex hypervectors, exact unbinding
fhrr = create_fhrr_model(dim=512)

# MAP: Real hypervectors, approximate unbinding
map_model = create_map_model(dim=512)

# Binary: Discrete hypervectors, exact unbinding
binary = create_binary_model(dim=10000, bipolar=True)

# Same interface for all models!
for model in [fhrr, map_model, binary]:
    memory = VSAMemory(model)
    memory.add("concept")
    vec = memory["concept"]

NEW: Clifford Operators (v1.1.0)

Exact, compositional, invertible transformations for reasoning:

from vsax import create_fhrr_model, VSAMemory
from vsax.operators import create_left_of, create_agent, create_patient
from vsax.similarity import cosine_similarity

model = create_fhrr_model(dim=512)
memory = VSAMemory(model)
memory.add_many(["cup", "plate", "dog", "cat", "chase"])

# Use pre-defined spatial operators
LEFT_OF = create_left_of(512)

# Encode: cup LEFT_OF plate
scene = model.opset.bundle(
    memory["cup"].vec,
    LEFT_OF.apply(memory["plate"]).vec
)

# Query with inverse operator
RIGHT_OF = LEFT_OF.inverse()
answer = LEFT_OF.inverse().apply(model.rep_cls(scene))
# Returns plate with high similarity!

# Use pre-defined semantic operators
AGENT = create_agent(512)
PATIENT = create_patient(512)

# Encode: "dog chases cat"
sentence = model.opset.bundle(
    AGENT.apply(memory["dog"]).vec,
    memory["chase"].vec,
    PATIENT.apply(memory["cat"]).vec
)

# Query: Who is the AGENT?
who = AGENT.inverse().apply(model.rep_cls(sentence))
similarity = cosine_similarity(who.vec, memory["dog"].vec)
print(f"AGENT is 'dog': {similarity:.3f}")  # High similarity!

Pre-defined operators (NEW in Phase 2):

• Spatial: create_left_of, create_right_of, create_above, create_below, create_in_front_of, create_behind, create_near, create_far
• Semantic: create_agent, create_patient, create_theme, create_experiencer, create_instrument, create_location, create_goal, create_source

Key features:

• ✅ Exact inversion: op.inverse().apply(op.apply(v)) recovers original (similarity > 0.999)
• ✅ Compositional: Combine operators algebraically with compose()
• ✅ Typed: Semantic metadata (SPATIAL, SEMANTIC, TEMPORAL, etc.)
• ✅ Reproducible: Same dimension always produces same operator
• ✅ FHRR-compatible: Phase-based transformations for complex hypervectors

NEW: Spatial Semantic Pointers (v1.2.0)

Continuous spatial encoding for 2D navigation and spatial reasoning:

from vsax import create_fhrr_model, VSAMemory
from vsax.spatial import SpatialSemanticPointers, SSPConfig

model = create_fhrr_model(dim=1024)
memory = VSAMemory(model)

# Create 2D spatial encoder
config = SSPConfig(dim=1024, num_axes=2, axis_names=["X", "Y"])
ssp = SpatialSemanticPointers(model, memory, config)

# Add objects to memory
memory.add_many(["table", "chair", "plant"])

# Bind objects to locations in a 10x10 room
table_at_loc = ssp.bind_object_location("table", [3.0, 7.0])
chair_at_loc = ssp.bind_object_location("chair", [5.0, 7.5])
plant_at_loc = ssp.bind_object_location("plant", [8.0, 3.0])

# Bundle into scene
from vsax.representations import ComplexHypervector
scene_vec = model.opset.bundle(
    table_at_loc.vec,
    model.opset.bundle(chair_at_loc.vec, plant_at_loc.vec)
)
scene = ComplexHypervector(scene_vec)

# Query: What's at location (3.0, 7.0)?
result = ssp.query_location(scene, [3.0, 7.0])
# Returns hypervector similar to "table"!

# Query: Where is the plant?
location_hv = ssp.query_object(scene, "plant")
decoded_location = ssp.decode_location(
    location_hv,
    search_range=[(0, 10), (0, 10)],
    resolution=20
)
# Returns approximately [8.0, 3.0]!

Key features:

• ✅ Continuous encoding: Positions encoded as X^x ⊗ Y^y using fractional powers
• ✅ What/Where queries: Both "what is at (x,y)" and "where is X" supported
• ✅ Scene shifting: Transform entire scenes spatially
• ✅ Similarity maps: Visualize spatial distributions

NEW: Vector Function Architecture (v1.2.0)

Function approximation in RKHS for density estimation, regression, and image processing:

from vsax import create_fhrr_model, VSAMemory
from vsax.vfa import DensityEstimator, NonlinearRegressor, ImageProcessor
import jax.numpy as jnp

model = create_fhrr_model(dim=2048)

# 1. Density Estimation
estimator = DensityEstimator(model, VSAMemory(model), bandwidth=0.5)
samples = jax.random.normal(jax.random.PRNGKey(42), (100,))
estimator.fit(samples)
density = estimator.evaluate(jnp.array([0.0, 1.0, 2.0]))

# 2. Nonlinear Regression
regressor = NonlinearRegressor(model, VSAMemory(model))
x_train = jnp.linspace(0, 2 * jnp.pi, 50)
y_train = jnp.sin(x_train)
regressor.fit(x_train, y_train)
y_pred = regressor.predict(jnp.array([0.5, 1.0, 1.5]))
r2_score = regressor.score(x_train, y_train)

# 3. Image Processing
processor = ImageProcessor(model, VSAMemory(model))
image = jnp.zeros((16, 16))
image = image.at[4:12, 4:12].set(1.0)  # White square
image_hv = processor.encode(image)
reconstructed = processor.decode(image_hv, shape=(16, 16))

# Blend two images
image2 = jnp.zeros((16, 16))
image2_hv = processor.encode(image2)
blended = processor.blend(image_hv, image2_hv, alpha=0.5)

Key features:

• ✅ Kernel density estimation: Estimate probability densities from samples
• ✅ Nonlinear regression: Fit complex functions with scikit-learn-like API
• ✅ Image encoding: Compact hypervector representation of images
• ✅ Function arithmetic: Add, scale, shift, and convolve functions
• ✅ RKHS approximation: Reproducing Kernel Hilbert Space function representation

Advanced: Manual Model Creation

You can still create models manually if you need custom configuration:

from vsax import VSAModel, ComplexHypervector, FHRROperations, sample_complex_random

model = VSAModel(
    dim=512,
    rep_cls=ComplexHypervector,
    opset=FHRROperations(),
    sampler=sample_complex_random
)

See docs/design-spec.md for complete technical specification.

Development Status

Currently in Iteration 6: I/O & Persistence ✅

Completed

Iteration 1 (v0.1.0): Foundation & Infrastructure ✅

• ✅ Core abstract classes (AbstractHypervector, AbstractOpSet)
• ✅ VSAModel dataclass
• ✅ Package structure
• ✅ Testing infrastructure (pytest, coverage)
• ✅ CI/CD pipeline (GitHub Actions)
• ✅ Documentation site (MkDocs)
• ✅ Development tooling (ruff, mypy)

Iteration 2 (v0.2.0): All 3 Representations + All 3 OpSets ✅

• ✅ ComplexHypervector, RealHypervector, BinaryHypervector
• ✅ FHRROperations, MAPOperations, BinaryOperations
• ✅ Sampling utilities (sample_random, sample_complex_random, sample_binary_random)
• ✅ 175 comprehensive tests with 96% coverage
• ✅ Full integration tests for all model combinations

Iteration 3 (v0.3.0): VSAMemory + Factory Functions ✅

• ✅ VSAMemory class - dictionary-style symbol management
• ✅ Factory functions (create_fhrr_model, create_map_model, create_binary_model)
• ✅ Utility functions (coerce_to_array, validation helpers)
• ✅ 230 tests with 89% coverage
• ✅ Comprehensive documentation guides

Iteration 4 (v0.4.0): Encoders + Examples ✅ FIRST USABLE RELEASE!

• ✅ ScalarEncoder - Numeric values with power encoding
• ✅ SequenceEncoder - Ordered sequences (lists, tuples)
• ✅ SetEncoder - Unordered collections (sets)
• ✅ DictEncoder - Key-value pairs (dictionaries)
• ✅ GraphEncoder - Graph structures (edge lists)
• ✅ AbstractEncoder - Base class for custom encoders
• ✅ Complete integration examples for all 3 models
• ✅ Custom encoder examples (DateEncoder, ColorEncoder)
• ✅ 280+ tests with 92%+ coverage

Iteration 5 (v0.5.0): Similarity Metrics & Utilities ✅

• ✅ Cosine, dot, and Hamming similarity functions
• ✅ Batch operations with JAX vmap (vmap_bind, vmap_bundle, vmap_similarity)
• ✅ Visualization utilities (pretty_repr, format_similarity_results)
• ✅ GPU-accelerated similarity search
• ✅ Comprehensive examples (similarity_search.py, batch_operations.py)
• ✅ 319 tests with 95%+ coverage

Iteration 6 (v0.6.0): I/O & Persistence ✅

• ✅ save_basis() and load_basis() functions
• ✅ JSON serialization for all 3 models
• ✅ Round-trip vector preservation
• ✅ Dimension and type validation
• ✅ Comprehensive tests (339 tests, 96% coverage)
• ✅ Complete examples and documentation

Coming Next

Iteration 7 (v1.0.0): Full Documentation & Production Release

• Complete API documentation
• Tutorial notebooks
• Production-ready v1.0.0 release

See todo.md for the complete development roadmap.

Documentation

• Full Documentation
• Getting Started
• Modeling Guide - 7-step workflow for building VSA applications
• Tutorials
  • MNIST Classification - Image classification with VSA
  • Knowledge Graph Reasoning - Multi-hop reasoning with relational facts
  • Kanerva's "Dollar of Mexico" - Analogical reasoning and mapping
  • Word Analogies & Random Indexing - Word embeddings and semantic similarity
  • Understanding VSA Models - Compare FHRR, MAP, and Binary models
  • VSA for Edge Computing - Lightweight alternative to neural networks
  • Hierarchical Structures - Trees and nested composition
  • Multi-Modal Concept Grounding - Fuse vision, language, and arithmetic
  • Neural-Symbolic Fusion (HD-Glue) - Glue neural networks with VSA ✨ NEW
• VSAX vs Other Libraries - Why VSAX? Feature comparison with torchhd, hdlib, PyBHV
• Design Specification
• API Reference
• Contributing

Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines.

Development

Run tests:

pytest

Run tests with coverage:

pytest --cov=vsax --cov-report=term-missing

Type checking:

mypy vsax

Linting:

ruff check vsax tests

Build documentation:

mkdocs serve

License

VSAX is released under the MIT License. See LICENSE for details.

Citation

If you use VSAX in your research, please cite:

@software{vsax2025,
  title = {VSAX: Vector Symbolic Algebra for JAX},
  author = {Sarathy, Vasanth},
  year = {2025},
  url = {https://github.com/vasanthsarathy/vsax},
  version = {1.2.0}
}

Acknowledgments

VSAX is built on JAX and inspired by the VSA/HDC research community.