PyTensor Más Allá de PyMC: Construyendo Inferencia LLM en Python

Python

PyTensor

MLX

Numba

LLM

GGUF

Gemma

Lo que le falta a PyTensor para inferencia LLM — y lo sencillo que es construirlo. Una exploración nativa en Python de grafos simbólicos, pesos GGUF, C, Numba, MLX, y el camino hacia un runtime LLM componible.

Autor/a

Carlos Trujillo

Fecha de publicación

12 de julio de 2026

Introducción

La mayoría de los profesionales conocen PyTensor a través de PyMC. Escribimos un modelo probabilístico, PyMC construye un grafo simbólico y PyTensor compila las matemáticas en algo que una máquina puede ejecutar.

Esa descripción es precisa. También está incompleta.

PyTensor es un compilador simbólico general de tensores. No sabe qué es un prior, y no requiere una verosimilitud. Es el compilador de grafos en el corazón de la programación probabilística, pero su arquitectura nunca estuvo limitada a ese dominio. Este artículo explora qué sucede cuando lo empujamos hacia algo que los autores probablemente no tenían en mente originalmente — inferencia LLM local — y lo que eso nos dice sobre hacia dónde podrían ir PyTensor y el proyecto pytensor-ml.

Esto es lo que PyTensor ya tiene: ejecución multi-backend (JAX, MLX, Numba, C) y optimización simbólica de grafos que simplifica el cálculo antes de llegar al hardware. La capa de ML vive un nivel arriba, en el ecosistema que lo rodea: pytensor-ml ya demuestra inferencia funcional para modelos pequeños, la descuantización GGUF viene de la biblioteca gguf, y el pipeline de Alchemize auto-genera módulos PyTensor a partir de metadatos GGUF.

Esto es lo que falta: cacheo KV de producción, batching continuo, carga mmap zero-copy y la optimización de rendimiento que hace que llama.cpp sirva modelos a escala. Esas brechas definen la hoja de ruta. Aquí construimos la rebanada vertical completa debajo de ellas: carga de modelo, ejecución simbólica, generación y validación numérica independiente.

La historia que cuenta este artículo no se trata de reemplazar a llama.cpp. Se trata de descubrir lo sencillo que es ensamblar un stack de inferencia LLM en PyTensor — pesos, tokenización, un transformador simbólico, un bucle de generación y validación numérica — una vez que el compilador de grafos se libera de sus suposiciones probabilísticas. Y se trata de lo que se vuelve posible cuando esas piezas se unen.

¿No sería genial poder ejecutar lo siguiente completamente en PyTensor? Y si piensas que no, imagina las ventajas:

Una definición, cualquier backend. Las mismas ecuaciones simbólicas se compilan a través de C, Numba, MLX o JAX. Cambiar el objetivo de hardware es un cambio de un argumento, no una reescritura.
Inferencia que compone. El modelo es un grafo dentro del ecosistema científico de Python — encadénalo con un posterior de PyMC, un optimizador de SciPy o cualquier cálculo personalizado, todo en el mismo framework.
Un grafo que puedes abrir. Cada operación es inspeccionable y reescribible. Puedes preguntar qué cambió una reescritura en vez de confiar en una caja negra.

from pathlib import Path
from pytensor import load_llm

model = load_llm(
    Path("/path/to/gemma-3n-E4B-it-lm-4bit/snapshot"),
    backend="c", # or any other.
)

result = inference(
    model,
    input="What's causal inference?",
    max_tokens=32,
)

Esa pequeña interfaz es nuestro destino.

Para llegar ahí, tenemos que ensamblar el stack completo nosotros mismos. PyTensor posee exactamente una capa — reescribir el grafo y enlazarlo a MLX, C/CVM o Numba — y todo alrededor es Python explícito: adaptadores de pesos que validan, mapean, descuantizan y orientan pesos GGUF o safetensors; adaptadores de tokenizador que aplican la plantilla de chat del modelo y producen IDs de token exactos; un modelo simbólico que expresa normalización, RoPE, atención, rutas residuales y MLPs; un runtime de generación que ejecuta prefill, actualiza el estado KV, elige tokens y se detiene; y una capa de reporte que devuelve texto, tiempos, memoria y verificaciones diferenciales. El primer paso en ese camino es Alchemize: lee los metadatos GGUF y nos entrega el mapa — el inventario de nombres de tensores, la estructura de bloques y los contratos exactos que faltan. Luego construimos — una pieza a la vez.

Construimos cada pieza contra Gemma 3n E4B a través de C/CVM, Numba y MLX, medimos lo que nos dice la velocidad actual y proyectamos hacia adelante.

El argumento

PyTensor ya tiene las abstracciones de grafo, reescritura y linker para convertirse en el núcleo computacional de un runtime LLM nativo en Python. Lo que falta no es el fundamento — es la ingeniería de producción encima de él. Y construir esa ingeniería en PyTensor es sorprendentemente sencillo.

¿Qué es Alchemize?

Alchemize es un transpilador autocorrectivo basado en LLM de PyMC Labs. Actúa como un agente de IA que compila modelos computacionales entre frameworks: PyMC, Stan, JAX, PyTorch y Rust — con validación numérica en cada paso. El agencia razona sobre el grafo computacional completo y aplica optimizaciones que haría un experto del dominio: fusión de bucles, pre-asignación de memoria, patrones de acceso amigables con el cache.

¡Lo que Alchemize revela!

Con el destino visible, podemos retroceder y seguir el camino que lo produjo — empezando por lo que falta.

Comenzamos con SmolLM2-135M-Instruct-Q4_K_M.gguf, un archivo GGUF de aproximadamente 105 MB. Nuestro primer intento es pedirle a Alchemize una implementación en PyTensor:

Show Alchemize call

from alchemize import compile_model

result = compile_model(
    model_path=Path("/path/to/SmolLM2-135M-Instruct-Q4_K_M.gguf"),
    target="pytensor",
)

Alchemize lee los metadatos GGUF y genera un módulo con el esqueleto de arquitectura correcto. Eso es valioso: nos da el mapa de nombres de tensores, el inventario de bloques y la estructura de capas sin escribir nada a mano.

Pero la implementación generada no puede ejecutar. Su suposición central de carga es incorrecta:

Show generated materialize_tensor

def materialize_tensor(tensor):
    data = getattr(tensor, "data", None)
    array = np.asarray(data)

    if np.issubdtype(array.dtype, np.floating):
        return np.asarray(array, dtype=np.float32)

    raise NotImplementedError("a dequantizer is required")

La mayoría de los tensores en este GGUF no son arreglos de punto flotante. Son valores cuantizados empaquetados. La implementación generada reconoce el problema y se detiene, pero nunca llama a gguf.dequantize.

La auditoría estática por lo tanto nos da:

Check	Result	Why
Provenance and syntax	pass	the artifact is pinned and valid Python
GGUF dequantization	`STATIC_FAIL`	packed weights are never dequantized
Attention head reshape	`STATIC_FAIL` audit flag	symbolic reshapes do not preserve the repaired runtime’s static contracts
Runtime	`BLOCKED`	generated code is never executed
Semantics	`UNVERIFIED`	valid syntax does not establish correct logits

Esto no es un fracaso. Es un mapa.

Alchemize aceleró el descubrimiento de arquitectura y nos dijo exactamente qué falta. El código generado nos da el esqueleto; la auditoría estática nos dice qué contratos necesitan reemplazos hechos a mano. Conservamos el mapa de nombres de tensores y el inventario de bloques, luego construimos las piezas que llenan los huecos: materialización, orientación, construcción del grafo, cacheo, ejecución y validación.

Construyendo inferencia de Gemma 3n en PyTensor

Una respuesta correcta depende de una cadena de contratos: descuantización de pesos cuantizados, orientación de tensores, tokenización, embeddings rotacionales, atención de consulta agrupada, cacheo KV y compilación del backend. Rompe un contrato y el modelo puede seguir produciendo texto plausible. Por eso construimos cada contrato explícitamente y los componemos en un pipeline funcional.

Nuestro objetivo es Gemma 3n E4B — un modelo mucho más grande y extraño que un típico primer experimento. Agrega cuatro flujos residuales AltUp, una ruta LAuReL de bajo rango aprendida, embeddings de token por capa, variantes de activación sparse y dense, atención deslizante y completa, y un vocabulario de 262.400 tokens. ¿Podemos expresar todo eso en el lenguaje simbólico de PyTensor sin cambiar el framework?

Component	E4B configuration
Decoder layers	35
Hidden size	2,048
Query / KV heads	8 / 2
Head dimension	256
MLP width	16,384
Vocabulary	262,400
AltUp streams	4
LAuReL rank	64
Sliding window	512
Final logit softcap	30

Este caso valida deliberadamente la generación de múltiples tokens con prefijo completo. No retiene un cache KV entre pasos de decodificación. En cambio, cada paso reconstruye el prefijo completo, crea la topología shared-KV del modelo para ese forward, y lo descarta después.

Transmite pesos en lugar de expandir el modelo

El checkpoint de 3.86 GB almacena cada módulo lineal affine-4 como pesos uint32 empaquetados más escalas y sesgos BF16. Ocho valores de cuatro bits ocupan una palabra:

W_{r,c} = q_{r,c}s_{r,g(c)} + b_{r,g(c)}.

Aquí, q_{r,c} es el nibble almacenado en la fila de salida r y la columna de entrada c; s_{r,g(c)} y b_{r,g(c)} son la escala y el sesgo para su grupo.

Expandir completamente todos los parámetros lógicos necesitaría aproximadamente 25.6 GiB solo para pesos FP32. En cambio, nosotros:

cargar solo las filas de embedding solicitadas,
descuantizar una capa de decodificador a la vez,
liberarla antes de cargar la siguiente capa, y
proyectar los logits de vocabulario en bloques de 4.096 filas.

Show weight streaming

from cetagostini.utils.pytensor.weights import Gemma3nWeightLoader

with Gemma3nWeightLoader.from_snapshot(snapshot_path) as loader:
    token_rows = loader.load_input_embedding_rows(token_ids)
    layer_0 = loader.load_layer(0)

La transmisión cambia el problema de “mantener el modelo expandido” a “mantener la capa actualmente expandida”.

Escribe las ecuaciones una sola vez

La normalización compartida es PyTensor ordinario:

Show rmsnorm_symbolic

from cetagostini.utils.pytensor.ops import rmsnorm_symbolic

normalized = rmsnorm_symbolic(hidden, gamma, eps=1e-6)

El detalle importante no es la fórmula. Es que rmsnorm_symbolic no sabe nada sobre C, Numba o MLX.

Lo mismo es cierto para la atención de consulta agrupada, RoPE, máscaras, AltUp y LAuReL. Las rutas GELU sparse y dense de Gemma producen dos FunctionGraphs especializados; la atención completa versus deslizante llega como datos de máscara y RoPE. El orden de operaciones sigue la implementación fijada de MLX-LM — incluyendo el aparente residual repetido de LAuReL y el patrón de dispersión GELU sparse leído del checkpoint.

Elige el backend en la compilación

La selección del backend es ahora una utilidad pequeña y reutilizable:

Show backend selection

from cetagostini.utils.pytensor.backends import get_mode

c_mode = get_mode("c")
numba_mode = get_mode("numba")
mlx_mode = get_mode("mlx")

c_layer = pytensor.function(layer_inputs, layer_output, mode=c_mode)
numba_layer = pytensor.function(layer_inputs, layer_output, mode=numba_mode)
mlx_layer = pytensor.function(layer_inputs, layer_output, mode=mlx_mode)

La definición del modelo no cambió. Solo cambiaron la política de linker y reescritura.

Un modelo, múltiples experimentos con el compilador

Si cambiamos una ecuación simbólica, cada backend la hereda. Si cambiamos solo una reescritura o linker, el modelo permanece fijo. Esa separación es la contribución principal de PyTensor a este experimento.

Ejecuta inferencia de Gemma desde Python

El mismo punto de entrada público ahora apunta a un artefacto y backend diferente:

from pathlib import Path
from cetagostini.utils.pytensor import Gemma3n, inference

gemma = Gemma3n.from_snapshot(
    Path("/path/to/gemma-3n-E4B-it-lm-4bit/snapshot"),
    backend="c",
)

result = inference(
    gemma,
    input="What's causal inference?",
    max_tokens=32,
)

result.output, result.output_token_ids

('## Causal Inference: Understanding "Why" vs. "Correlation"\n\n'
 'Causal inference is a branch of statistics and statistics is a field '
 'of statistics that aims',
 (1408, 565, 90718, 157036, 236787, 40411, 623, 11355, 236775,
  7728, 236761, 623, 131426, 236775, 108, 236780, 90718, 34711,
  563, 496, 9911, 529, 14906, 532, 14906, 563, 496, 2135,
  529, 14906, 600, 17269))

Esa es una continuación real, no una sola predicción de siguiente token. Tampoco es una prosa pulida: la decodificación greedy alcanza el límite de 32 tokens a mitad de oración y se vuelve repetitiva después de que la ruta diferencial se separa. El punto es hacer la generación inspeccionable, no presentar un benchmark de calidad lingüística.

Show validation report

result.report.validation

{
    'final_top1_match': True,
    'all_top1_match': False,
    'pearson_mean': 0.998932415848074,
    'top10_overlap_mean': 9.533333333333333,
    'thresholds_passed': True,
    'scope': 'all_generated_prefixes_fresh_cache',
    'all_generated_tokens_match': False,
}

El prompt inicial de 15 tokens pasa los umbrales de publicación y selecciona el mismo siguiente token en PyTensor y MLX-LM. La generación entonces valida cada prefijo creciente de PyTensor contra un forward de MLX-LM sobre ese prefijo exacto. Cada llamada al oráculo crea un objeto shared-KV fresco y lo descarta; ningún estado de decodificación sobrevive al siguiente paso.

Las primeras 20 decisiones generadas coinciden. En el token 21, PyTensor elige 9911 mientras el oráculo MLX de prefijo fresco elige 2135. El reporte registra esa divergencia y sigue evaluando MLX en el prefijo de PyTensor posterior, por lo que las comparaciones posteriores permanecen bien definidas. No compara dos flujos derivando independientemente, y ya no descarta una generación completada costosa porque stream_generate cacheado sigue una ruta diferente.

Multi-token result	Value
Visible tokens	32
Matching fresh-prefix decisions	31 / 32
First divergence	token 21
Prefix graph compilations	32
Additional generation wall time	1,612.038 s
Peak process RSS	9,668.61 MiB
Stop reason	`max_tokens`

El costo es la lección. Después del paso inicial validado del prompt, el bucle restante de generación de prefijo completo tomó 1.612,038 segundos, o aproximadamente 26,9 minutos. Esta es una estrategia de ejecución educativa que expone cada grafo y comparación; no es un decodificador eficiente.

C/CVM vs Numba vs MLX: comparación de rendimiento

Los tres backends compilan las mismas ecuaciones simbólicas y seleccionan el mismo token top-1 que el oráculo independiente de MLX-LM en las 20 posiciones del prompt. Su correlación de Pearson media de todos los logits con el oráculo es 0.9986, y su superposición media del top-10 es 9.7 de 10.

Metric	C/CVM	Numba	MLX
Graph compilation	7.738 s	0.775 s	0.755 s
One full forward (20 positions)	51.441 s	53.800 s	50.211 s
Mean per-layer time	1.272 s	1.266 s	1.246 s
Logit projection	3.024 s	4.700 s	3.086 s
Whole-process peak RSS	4,523 MiB	4,663 MiB	5,111 MiB

Estas son mediciones de una sola ejecución de un pipeline educativo transmitido intencionalmente, no un benchmark de throughput. MLX compila aproximadamente 10 veces más rápido que C y registra el forward más corto, pero solo por un 2.4%. C todavía gana la proyección del vocabulario de 262.400. La descuantización de pesos, la orquestación en Python y la transmisión por capa dominan lo suficiente como para que los tres forwards completos permanezcan en el mismo rango.

Cómo la sonda MLX fallida se convirtió en un backend funcional

La primera sonda MLX falló porque el borrador generado dependía de formas y operaciones que el linker fijado no podía reducir de forma segura. No parcheamos el paquete PyTensor instalado ni mutamos un registro de despacho global de proceso. En su lugar, expresamos las proyecciones como multiplicaciones de matriz de rango 2 con reshapes explícitos, mantuvimos la atención en primitivas ordinarias de tensores, y reemplazamos los sitios de clip de AltUp de Gemma con un ayudante simbólico local del repositorio construido con comparaciones y where.

El grafo resultante se compila a través de pytensor.compile.mode.MLX integrado en PyTensor 3.1.2. MLX evalúa de forma lazy, así que el runner materializa 103 límites ordenados: dos proyecciones iniciales, 35 capas de decodificador, un desincrustar final y 65 bloques de vocabulario. Los tensores intermedios permanecen residentes en el dispositivo; solo los bloques de logits completados cruzan de vuelta a arreglos NumPy propietarios. El pico medido del asignador MLX fue 455 MiB por encima de su línea base cercana a cero, mientras que el RSS de todo el proceso alcanzó 5.111 MiB.

Los reportes completos están disponibles para el oráculo independiente de MLX-LM, C/CVM, Numba y MLX. Un validador separado recargó los logits crudos temporales, recalculó cada métrica en vez de confiar en los reportes, y pasó 896 de 896 pruebas. Los arreglos crudos grandes no están comprometidos; los reportes preservan sus formas, dtypes, conteos de bytes y hashes criptográficos. La ejecución anterior de generación de 32 tokens permanece en gemma3n_pytensor_generation.json.

Lo que nos dice la velocidad actual

Los números anteriores no son benchmarks. Son mediciones de un pipeline educativo, y deberían leerse así. Pero son honestos, y nos dicen algo preciso sobre hacia dónde necesita dirigirse el esfuerzo de ingeniería.

Gemma 3n a través de regeneración de prefijo completo funciona a aproximadamente 0.02 tokens por segundo. Ese es el costo de recompilar y reevaluar el prefijo completo para cada token generado. Es deliberadamente costoso — valida corrección — pero no es como funcionaría un sistema de producción.

La hoja de ruta de ingeniería desde aquí es clara:

Bottleneck	Current state	What unlocks it
KV cache	fixed-capacity, O(C) write per layer	paged or ring-buffer cache, continuous batching
Weight loading	per-layer streaming (dequantize on demand)	mmap zero-copy, quantized kernels
Backend execution	C/CVM, Numba, and MLX on Apple Silicon	backend-native quantized kernels and less host-device movement
Graph compilation	recompilation per prefix length	cached compiled functions per shape, or JIT

llama.cpp ha dedicado años a cada fila de esa tabla. PyTensor tiene el compilador de grafos y la arquitectura multi-backend; todavía no tiene la infraestructura de serving. La pregunta no es si PyTensor puede igualar el throughput de llama.cpp hoy — no puede — sino si las piezas están en su lugar para construir esa infraestructura en Python. La respuesta, después de este experimento, es sí.

La ventaja de composabilidad

llama.cpp es un motor de inferencia autocontenido. PyTensor es un compilador de grafos que puede componerse con transformaciones de JAX, operaciones de NumPy, optimizadores de SciPy y el resto del stack científico de Python. En el momento en que la inferencia LLM vive dentro de ese ecosistema, puedes encadenarla con análisis Bayesiano, optimización basada en gradientes o cálculo simbólico personalizado — flujos de trabajo para los que un motor C++ independiente nunca fue diseñado para soportar.

Donde PyTensor y llama.cpp difieren — y por qué importa

Ahora podemos hacer la comparación precisa:

Capability	This PyTensor stack	`llama.cpp`
Inspectable symbolic graph	yes	not its primary user abstraction
User-defined graph rewrites	yes	no equivalent Python rewrite database
C, Numba, and MLX experiments	demonstrated on Gemma 3n E4B	purpose-built CPU, Metal, CUDA, and other backends
GGUF parsing	supplied by `gguf-py` adapter	built in
Native quantized matmul	not implemented here	built in
Tokenization and chat templates	Python adapters	built in
Autoregressive loop	Python, model-specific	built in
Production KV cache and batching	no	built in
Broad architecture support	one validated fixture (Gemma 3n)	broad, maintained model coverage
Composability with scientific Python	native	not designed for it

llama.cpp gana en ingeniería de inferencia integrada. Eso es exactamente para lo que fue construido.

Pero el ecosistema local de LLM no se trata solo de ejecutar un modelo rápido. Se trata de lo que haces con el modelo después.

La historia reciente de Ollama — documentada extensamente por sleepingrobots — muestra qué sucede cuando un wrapper oscurece sus dependencias y pivota a la nube. El ecosistema necesita alternativas construidas sobre fundamentos honestos y transparentes. PyTensor y llama.cpp son esos fundamentos.

llama.cpp es el motor en C++ para ejecutar cualquier modelo GGUF rápido. Posee todo el stack de serving: kernels cuantizados, gestión de cache KV, batching continuo y soporte amplio de arquitecturas. Si necesitas servir modelos a escala de producción hoy, llama.cpp es la respuesta.
PyTensor es el compilador de grafos nativo en Python para estudiar, modificar y componer inferencia LLM. No solo ejecutas un modelo — puedes preguntar qué cambió una reescritura, compilar las mismas ecuaciones a través de otro linker, encadenar la inferencia con un posterior Bayesiano o un bucle de optimización personalizado, y validar cada contrato numéricamente.

La afirmación honesta no es “PyTensor reemplaza a llama.cpp”. Es esta:

PyTensor más adaptadores explícitos en Python es una alternativa nativa en Python creíble a llama.cpp que compone con el stack de computación científica — análisis Bayesiano, optimización, grafos de cálculo personalizados — de maneras para las que un motor C++ independiente nunca fue diseñado.

Si alguien construye cacheo KV, carga mmap zero-copy y batching continuo sobre el backend JAX de PyTensor, y Alchemize madura para auto-generar nuevas arquitecturas a partir de metadatos GGUF, podríamos ver a PyTensor convertirse en el “lenguaje” para ejecutar múltiples LLMs como lo es llama.cpp hoy. No más rápido en inferencia — o quizás en algún momento — pero más capaz como un runtime LLM de propósito general que se conecta a flujos de trabajo para los que llama.cpp nunca fue diseñado.

Conclusiones

PyTensor es un compilador de grafos general. PyMC es su consumidor más visible, no el límite de lo que puede expresar.
Construir inferencia LLM en PyTensor es sencillo. Cargadores de pesos, transformadores simbólicos, caches KV, bucles de generación y validación numérica — todo ensamblado desde Python sin modificar el framework.
Una definición simbólica sobrevive a múltiples backends. Gemma 3n se ejecuta a través de C/CVM, Numba y MLX sin una segunda implementación del modelo.
Los logits independientes superan a la prosa plausible. Tokens exactos y concordancia numérica en todas las posiciones son evidencia más fuerte que un texto que simplemente parece correcto.
PyTensor y llama.cpp son complementarios. llama.cpp domina el serving en producción. PyTensor domina la transparencia del grafo y la composabilidad con Python científico.
Las brechas restantes son de ingeniería, no de arquitectura. Kernels cuantizados nativos, cacheo KV paginado, carga mmap zero-copy y cobertura más amplia de backends son todo construible sobre el fundamento existente de PyTensor. El proyecto pytensor-ml ya ha comenzado este trabajo.

Lecturas recomendadas:

El experimento registrado utilizó:

Component	Version or pin
Python	`3.13.14`
PyTensor	`3.1.2`
PyTensor-ML	`f6ecf81d58da180cce50b77a43cf5d2c3d95e470`
MLX	`0.32.0`
Numba	`0.65.1`
MLX-LM	`0.31.3`
Machine	Apple M3 Max, 128 GB unified memory

El entorno exacto está registrado en environment.yml. Las celdas de código no se ejecutan durante la compilación del sitio web porque los artefactos del modelo están intencionalmente excluidos de Git; los reportes JSON comprometidos preservan los hashes de los artefactos y las salidas medidas. La suite de PyTensor pasó 934 pruebas con 30 omisiones, y las pruebas del runner con modelo real pasaron las 189 verificaciones; el registro compacto está en gemma3n_test_summary.json.

El paquete público, las implementaciones específicas del modelo, las pruebas, el borrador generado, las auditorías y los reportes de resultados sanitizados están disponibles en el repositorio del sitio.

--- title: "PyTensor Beyond PyMC: Building LLM Inference in Python" pagetitle: "PyTensor Beyond PyMC: Building LLM Inference in Python — Carlos Trujillo" author: "Carlos Trujillo" date: "2026-07-12" description: "What PyTensor is missing for LLM inference—and how straightforward it is to build. A Python-native exploration of symbolic graphs, GGUF weights, C, Numba, MLX, and the path to a composable LLM runtime." aliases: - gemma3n_e4b_symbolic_pytensor_inference.html categories: [python, pytensor, mlx, numba, llm, gguf, gemma] image: "/images/alchemize_pytensor_mlx_gemma_3n.jpg" image-alt: "Alchemize logo over PyTensor and MLX symbols representing LLM inference in Python" jupyter: alchemize_pytensor_mlx_gemma_3n execute: eval: false freeze: false format: html: code-fold: false code-tools: true code-overflow: wrap --- # Introduction Most practitioners meet [PyTensor](https://pytensor.readthedocs.io/) through PyMC. We write a probabilistic model, PyMC builds a symbolic graph, and PyTensor compiles the mathematics into something a machine can execute. That description is accurate. It is also incomplete. PyTensor is a general symbolic tensor compiler. It does not know what a prior is, and it does not require a likelihood. It is the graph compiler at the heart of probabilistic programming, but its architecture was never limited to that domain. This article explores what happens when we push it toward something the authors may not have originally intended—local LLM inference—and what that tells us about where PyTensor, and the [pytensor-ml](https://github.com/pymc-labs/pytensor-ml) project, could go next. Here is what PyTensor itself already has: multi-backend execution (JAX, MLX, Numba, C) and symbolic graph optimization that simplifies computation before it hits hardware. The ML layer lives one level up, in the ecosystem around it: [`pytensor-ml`](https://github.com/jessegrabowski/pytensor_ml) already demonstrates working inference for small models, GGUF dequantization comes from the [`gguf`](https://github.com/ggml-org/llama.cpp/tree/master/gguf-py) library, and the [Alchemize](https://github.com/pymc-labs/alchemize) pipeline auto-generates PyTensor modules from GGUF metadata. Here is what is missing: production KV caching, continuous batching, mmap zero-copy loading, and the performance tuning that makes `llama.cpp` serve models at scale. Those gaps define the roadmap. Here we build the complete vertical slice underneath them: model loading, symbolic execution, generation, and independent numerical validation. The story this article tells is not about replacing `llama.cpp`. It is about discovering how straightforward it is to assemble an LLM inference stack in PyTensor—weights, tokenization, a symbolic transformer, a generation loop, and numerical validation—once the graph compiler is freed from its probabilistic assumptions. And it is about what becomes possible when those pieces come together. It will be really cool be able to run the following in full pytensor no? And if you think it would not, imagine the advantages: - **One definition, any backend.** The same symbolic equations compile through C, Numba, MLX, or JAX. Swapping hardware targets is a one-argument change, not a rewrite. - **Inference that composes.** The model is a graph inside the scientific Python stack—chain it with a PyMC posterior, a SciPy optimizer, or any custom computation, all in the same framework. - **A graph you can open.** Every operation is inspectable and rewritable. You can ask what a rewrite changed instead of trusting a black box. ```python from pathlib import Path from pytensor import load_llm model = load_llm( Path("/path/to/gemma-3n-E4B-it-lm-4bit/snapshot"), backend="c", # or any other. ) result = inference( model, input="What's causal inference?", max_tokens=32, ) ``` That small interface is our destination. To reach it, we have to assemble the full stack ourselves. PyTensor owns exactly one layer—rewriting the graph and linking it to MLX, C/CVM, or Numba—and everything around it is explicit Python: weight adapters that validate, map, dequantize, and orient GGUF or safetensors weights; tokenizer adapters that apply the model's chat template and produce exact token IDs; a symbolic model that expresses normalization, RoPE, attention, residual paths, and MLPs; a generation runtime that runs prefill, updates KV state, chooses tokens, and stops; and a report layer that returns text, timings, memory, and differential checks. The first step on that path is [Alchemize](https://github.com/pymc-labs/alchemize): it reads the GGUF metadata and hands us the map—the tensor-name inventory, the block structure, and the exact contracts that are missing. Then we build—one piece at a time. We build each piece against Gemma 3n E4B through C/CVM, Numba, and MLX, measure what the current speed tells us, and project forward. ::: {.callout-note collapse="true" appearance="simple"} ## The argument PyTensor already has the graph, rewrite, and linker abstractions to become the computational core of a Python-native LLM runtime. What is missing is not the foundation—it is the production engineering on top of it. And building that engineering in PyTensor is surprisingly straightforward. ::: ::: {.callout-note collapse="true" appearance="simple"} ## What is Alchemize? [Alchemize](https://github.com/pymc-labs/alchemize) is an LLM-based, self-correcting transpiler from PyMC Labs. It acts as an AI agent that compiles computational models between frameworks: PyMC, Stan, JAX, PyTorch, and Rust—with numerical validation at every step. The agent reasons about the full computational graph and applies optimizations a domain expert would: loop fusion, memory pre-allocation, cache-friendly access patterns. ::: # What Alchemize reveals! With the destination visible, we can rewind and follow the path that produced it—starting with what is missing. We begin with `SmolLM2-135M-Instruct-Q4_K_M.gguf`, a roughly 105 MB GGUF file. Our first attempt is to ask [Alchemize](https://github.com/pymc-labs/alchemize) for a PyTensor implementation: ```{python} #| code-fold: true #| code-summary: "Show Alchemize call" from alchemize import compile_model result = compile_model( model_path=Path("/path/to/SmolLM2-135M-Instruct-Q4_K_M.gguf"), target="pytensor", ) ``` Alchemize reads the GGUF metadata and generates a module with the right architecture skeleton. That is valuable: it gives us the tensor-name map, the block inventory, and the layer structure without writing any of it by hand. But the generated implementation cannot run. Its central loading assumption is wrong: ```{python} #| code-fold: true #| code-summary: "Show generated materialize_tensor" def materialize_tensor(tensor): data = getattr(tensor, "data", None) array = np.asarray(data) if np.issubdtype(array.dtype, np.floating): return np.asarray(array, dtype=np.float32) raise NotImplementedError("a dequantizer is required") ``` Most tensors in this GGUF are not floating arrays. They are packed quantized values. The generated implementation recognizes the problem and stops, but it never calls `gguf.dequantize`. The static audit therefore gives us: | Check | Result | Why | |---|---|---| | Provenance and syntax | pass | the artifact is pinned and valid Python | | GGUF dequantization | `STATIC_FAIL` | packed weights are never dequantized | | Attention head reshape | `STATIC_FAIL` audit flag | symbolic reshapes do not preserve the repaired runtime's static contracts | | Runtime | `BLOCKED` | generated code is never executed | | Semantics | `UNVERIFIED` | valid syntax does not establish correct logits | This is not a failure. It is a map. Alchemize accelerated architecture discovery and told us exactly what is missing. The generated code gives us the skeleton; the static audit tells us which contracts need hand-built replacements. We keep the tensor-name map and block inventory, then build the pieces that fill the gaps: materialization, orientation, graph construction, caching, execution, and validation. # Building Gemma 3n inference in PyTensor A correct answer depends on a chain of contracts: quantized weight dequantization, tensor orientation, tokenization, rotary embeddings, grouped-query attention, KV caching, and backend compilation. Break one contract and the model may still produce plausible text. So we build each contract explicitly and compose them into a working pipeline. Our target is Gemma 3n E4B—a much larger and stranger model than a typical first experiment. It adds four AltUp residual streams, a learned low-rank LAuReL path, per-layer token embeddings, sparse and dense activation variants, sliding and full attention, and a 262,400-token vocabulary. Can we express all of that in PyTensor's symbolic language without changing the framework? | Component | E4B configuration | |---|---:| | Decoder layers | 35 | | Hidden size | 2,048 | | Query / KV heads | 8 / 2 | | Head dimension | 256 | | MLP width | 16,384 | | Vocabulary | 262,400 | | AltUp streams | 4 | | LAuReL rank | 64 | | Sliding window | 512 | | Final logit softcap | 30 | This case deliberately validates **multi-token, full-prefix generation**. It does not retain a KV cache between decoding steps. Instead, every step rebuilds the complete prefix, creates the model's shared-KV topology for that one forward, and discards it afterward. ## Stream weights instead of expanding the model The 3.86 GB checkpoint stores each affine-4 linear module as packed `uint32` weights plus BF16 scales and biases. Eight four-bit values occupy one word: $$ W_{r,c} = q_{r,c}s_{r,g(c)} + b_{r,g(c)}. $$ Here, $q_{r,c}$ is the stored nibble at output row $r$ and input column $c$; $s_{r,g(c)}$ and $b_{r,g(c)}$ are the scale and bias for its group. Fully expanding all logical parameters would need about **25.6 GiB** for FP32 weights alone. We instead: - load only requested embedding rows, - dequantize one decoder layer at a time, - release it before loading the next layer, and - project vocabulary logits in chunks of 4,096 rows. ```{python} #| code-fold: true #| code-summary: "Show weight streaming" from cetagostini.utils.pytensor.weights import Gemma3nWeightLoader with Gemma3nWeightLoader.from_snapshot(snapshot_path) as loader: token_rows = loader.load_input_embedding_rows(token_ids) layer_0 = loader.load_layer(0) ``` Streaming changes the problem from “hold the expanded model” to “hold the current expanded layer.” ## Write the equations once The shared normalization is ordinary PyTensor: ```{python} #| code-fold: true #| code-summary: "Show rmsnorm_symbolic" from cetagostini.utils.pytensor.ops import rmsnorm_symbolic normalized = rmsnorm_symbolic(hidden, gamma, eps=1e-6) ``` The important detail is not the formula. It is that `rmsnorm_symbolic` knows nothing about C, Numba, or MLX. The same is true for grouped-query attention, RoPE, masks, AltUp, and LAuReL. Gemma's sparse and dense GELU paths produce two specialized `FunctionGraph`s; full versus sliding attention arrives as mask and RoPE data. The operation order follows the pinned MLX-LM implementation—including LAuReL's apparent repeated residual and the sparse GELU sparsity pattern read from the checkpoint. ## Choose the backend at compilation Backend selection is now a small, reusable utility: ```{python} #| code-fold: true #| code-summary: "Show backend selection" from cetagostini.utils.pytensor.backends import get_mode c_mode = get_mode("c") numba_mode = get_mode("numba") mlx_mode = get_mode("mlx") c_layer = pytensor.function(layer_inputs, layer_output, mode=c_mode) numba_layer = pytensor.function(layer_inputs, layer_output, mode=numba_mode) mlx_layer = pytensor.function(layer_inputs, layer_output, mode=mlx_mode) ``` The model definition did not change. Only the linker and rewrite policy changed. ::: {.callout-tip appearance="simple"} ## One model, multiple compiler experiments If we change a symbolic equation, every backend inherits it. If we change only a rewrite or linker, the model remains fixed. That separation is PyTensor's main contribution to this experiment. ::: ## Run Gemma inference from Python The same public entry point now targets a different artifact and backend: ```{python} #| echo: true from pathlib import Path from cetagostini.utils.pytensor import Gemma3n, inference gemma = Gemma3n.from_snapshot( Path("/path/to/gemma-3n-E4B-it-lm-4bit/snapshot"), backend="c", ) result = inference( gemma, input="What's causal inference?", max_tokens=32, ) result.output, result.output_token_ids ``` ```text ('## Causal Inference: Understanding "Why" vs. "Correlation"\n\n' 'Causal inference is a branch of statistics and statistics is a field ' 'of statistics that aims', (1408, 565, 90718, 157036, 236787, 40411, 623, 11355, 236775, 7728, 236761, 623, 131426, 236775, 108, 236780, 90718, 34711, 563, 496, 9911, 529, 14906, 532, 14906, 563, 496, 2135, 529, 14906, 600, 17269)) ``` That is an actual continuation, not one next-token prediction. It is also not polished prose: greedy decoding reaches the 32-token cap mid-sentence and becomes repetitive after the differential path separates. The point is to make generation inspectable, not to present a language-quality benchmark. ```{python} #| code-fold: true #| code-summary: "Show validation report" result.report.validation ``` ```text { 'final_top1_match': True, 'all_top1_match': False, 'pearson_mean': 0.998932415848074, 'top10_overlap_mean': 9.533333333333333, 'thresholds_passed': True, 'scope': 'all_generated_prefixes_fresh_cache', 'all_generated_tokens_match': False, } ``` The initial 15-token prompt passes the publication thresholds and selects the same next token in PyTensor and MLX-LM. Generation then validates each growing PyTensor prefix against an MLX-LM forward over that exact prefix. Each oracle call creates a fresh shared-KV object and discards it; no decode state survives into the next step. The first 20 generated decisions agree. At token 21, PyTensor chooses `9911` while the fresh-prefix MLX oracle chooses `2135`. The report records that divergence and keeps evaluating MLX on the subsequent **PyTensor prefix**, so later comparisons remain well-defined. It does not compare two independently drifting streams, and it no longer throws away an expensive completed generation because cached `stream_generate` follows a different path. | Multi-token result | Value | |---|---:| | Visible tokens | 32 | | Matching fresh-prefix decisions | 31 / 32 | | First divergence | token 21 | | Prefix graph compilations | 32 | | Additional generation wall time | 1,612.038 s | | Peak process RSS | 9,668.61 MiB | | Stop reason | `max_tokens` | The cost is the lesson. After the initial validated prompt pass, the remaining full-prefix generation loop took **1,612.038 seconds**, or about **26.9 minutes**. This is an educational execution strategy that exposes every graph and comparison; it is not an efficient decoder. ### C/CVM vs Numba vs MLX: performance comparison All three backends compile the same symbolic equations and select the same top-1 token as the independent MLX-LM oracle at all 20 prompt positions. Their mean all-logit Pearson correlation with the oracle is 0.9986, and their mean top-10 overlap is 9.7 out of 10. | Metric | C/CVM | Numba | MLX | |---|---:|---:|---:| | Graph compilation | 7.738 s | 0.775 s | **0.755 s** | | One full forward (20 positions) | 51.441 s | 53.800 s | **50.211 s** | | Mean per-layer time | 1.272 s | 1.266 s | **1.246 s** | | Logit projection | **3.024 s** | 4.700 s | 3.086 s | | Whole-process peak RSS | **4,523 MiB** | 4,663 MiB | 5,111 MiB | These are single-run measurements of an intentionally streamed educational pipeline, not a throughput benchmark. MLX compiles about 10 times faster than C and records the shortest forward, but only by 2.4%. C still wins the 262,400-vocabulary projection. Weight dequantization, Python orchestration, and per-layer streaming dominate enough that the three full forwards remain in the same range. ::: {.callout-note collapse="true" appearance="simple"} ## How the failed MLX probe became a working backend The first MLX probe failed because the generated draft relied on shapes and operations that the pinned linker could not lower safely. We did not patch the installed PyTensor package or mutate a process-global dispatch registry. Instead, we expressed projections as rank-2 matrix multiplies with explicit reshapes, kept attention in ordinary tensor primitives, and replaced Gemma's AltUp clip sites with a repository-local symbolic helper built from comparisons and `where`. The resulting graph compiles through PyTensor 3.1.2's built-in `pytensor.compile.mode.MLX`. MLX evaluates lazily, so the runner materializes 103 ordered boundaries: two initial projections, 35 decoder layers, one final unembed, and 65 vocabulary chunks. Intermediate tensors remain device-resident; only completed logit chunks cross back to owning NumPy arrays. The measured MLX allocator peak was 455 MiB beyond its near-zero baseline, while whole-process RSS reached 5,111 MiB. ::: The complete reports are available for the independent [MLX-LM oracle](results/gemma3n_mlx_lm_oracle.json), [C/CVM](results/gemma3n_pytensor_c.json), [Numba](results/gemma3n_pytensor_numba.json), and [MLX](results/gemma3n_pytensor_mlx.json). A separate validator reloaded the temporary raw logits, recomputed every metric instead of trusting the reports, and passed [896 of 896 gates](results/gemma3n_report_validation.json). The large raw arrays are not committed; the reports preserve their shapes, dtypes, byte counts, and cryptographic hashes. The earlier 32-token generation run remains in [`gemma3n_pytensor_generation.json`](results/gemma3n_pytensor_generation.json). # What the current speed tells us The numbers above are not benchmarks. They are measurements of an educational pipeline, and they should be read that way. But they are honest, and they tell us something precise about where the engineering effort needs to go. Gemma 3n through full-prefix regeneration runs at roughly **0.02 tokens per second**. That is the cost of recompiling and re-evaluating the entire prefix for every generated token. It is deliberately expensive—it validates correctness—but it is not how a production system would work. The engineering roadmap from here is clear: | Bottleneck | Current state | What unlocks it | |---|---|---| | KV cache | fixed-capacity, O(C) write per layer | paged or ring-buffer cache, continuous batching | | Weight loading | per-layer streaming (dequantize on demand) | mmap zero-copy, quantized kernels | | Backend execution | C/CVM, Numba, and MLX on Apple Silicon | backend-native quantized kernels and less host-device movement | | Graph compilation | recompilation per prefix length | cached compiled functions per shape, or JIT | `llama.cpp` has spent years on every row of that table. PyTensor has the graph compiler and the multi-backend architecture; it does not yet have the serving infrastructure. The question is not whether PyTensor can match `llama.cpp`'s throughput today—it cannot—but whether the pieces are in place to build that infrastructure in Python. The answer, after this experiment, is yes. ::: {.callout-tip appearance="simple"} ## The composability advantage `llama.cpp` is a self-contained inference engine. PyTensor is a graph compiler that can compose with JAX transformations, NumPy operations, SciPy optimizers, and the rest of the scientific Python stack. The moment LLM inference lives inside that ecosystem, you can chain it with Bayesian analysis, gradient-based optimization, or custom symbolic computation—workflows that a standalone C++ engine was never designed to support. ::: # Where PyTensor and llama.cpp differ — and why that matters Now we can make the comparison precise: | Capability | This PyTensor stack | `llama.cpp` | |---|---|---| | Inspectable symbolic graph | yes | not its primary user abstraction | | User-defined graph rewrites | yes | no equivalent Python rewrite database | | C, Numba, and MLX experiments | demonstrated on Gemma 3n E4B | purpose-built CPU, Metal, CUDA, and other backends | | GGUF parsing | supplied by `gguf-py` adapter | built in | | Native quantized matmul | not implemented here | built in | | Tokenization and chat templates | Python adapters | built in | | Autoregressive loop | Python, model-specific | built in | | Production KV cache and batching | no | built in | | Broad architecture support | one validated fixture (Gemma 3n) | broad, maintained model coverage | | Composability with scientific Python | native | not designed for it | `llama.cpp` wins on integrated inference engineering. That is exactly what it is built for. But the local LLM ecosystem is not just about running one model fast. It is about what you do with the model afterward. The recent history of Ollama—[documented thoroughly by sleepingrobots](https://sleepingrobots.com/dreams/stop-using-ollama/)—shows what happens when a wrapper obscures its dependencies and pivots to cloud. The ecosystem needs alternatives built on honest, transparent foundations. PyTensor and `llama.cpp` are those foundations. - **`llama.cpp`** is the C++ engine for running any GGUF model fast. It owns the full serving stack: quantized kernels, KV cache management, continuous batching, broad architecture support. If you need to serve models at production scale today, `llama.cpp` is the answer. - **PyTensor** is the Python-native graph compiler for studying, modifying, and composing LLM inference. You do not just run a model—you can ask what a rewrite changed, compile the same equations through another linker, chain inference with a Bayesian posterior or a custom optimization loop, and validate every contract numerically. The honest claim is not "PyTensor replaces `llama.cpp`". It is this: [PyTensor]{style="color:var(--green-strong)"} plus explicit Python adapters is a credible *Python-native alternative* to `llama.cpp` that **composes with the scientific computing stack**—Bayesian analysis, optimization, custom computation graphs—in ways a standalone C++ engine was never designed for. If someone builds KV caching, mmap zero-copy loading, and continuous batching on PyTensor's JAX backend, and [Alchemize](https://github.com/pymc-labs/alchemize) matures for auto-generating new architectures from GGUF metadata, we could see PyTensor become the "language" for running multiple LLMs the way `llama.cpp` is today. Not faster at inference—or maybe at some point—but more capable as a general-purpose LLM runtime that plugs into workflows `llama.cpp` was never built for. # Conclusions 1. **PyTensor is a general graph compiler.** PyMC is its most visible consumer, not the boundary of what it can express. 2. **Building LLM inference in PyTensor is straightforward.** Weight loaders, symbolic transformers, KV caches, generation loops, and numerical validation—all assembled from Python without modifying the framework. 3. **One symbolic definition survives multiple backends.** Gemma 3n executes through C/CVM, Numba, and MLX without a second model implementation. 4. **Independent logits beat plausible prose.** Exact tokens and all-position numerical agreement are stronger evidence than plausible-looking text. 5. **PyTensor and `llama.cpp` are complementary.** `llama.cpp` owns production serving. PyTensor owns graph transparency and composability with scientific Python. 6. **The remaining gaps are engineering, not architecture.** Native quantized kernels, paged KV caching, mmap zero-copy loading, and broader backend coverage are all buildable on PyTensor's existing foundation. The pytensor-ml project has already begun this work. Recommended readings: 1. [PyTensor documentation](https://pytensor.readthedocs.io/) 2. [PyTensor graph rewriting](https://pytensor.readthedocs.io/en/latest/extending/graph_rewriting.html) 3. [PyTensor compilation modes](https://pytensor.readthedocs.io/en/latest/library/compile/mode.html) 4. [llama.cpp](https://github.com/ggml-org/llama.cpp) 5. [MLX-LM](https://github.com/ml-explore/mlx-lm) 6. [Alchemize](https://github.com/pymc-labs/alchemize) 7. [pytensor-ml](https://github.com/pymc-labs/pytensor-ml) 8. [Friends Don't Let Friends Use Ollama](https://sleepingrobots.com/dreams/stop-using-ollama/) The recorded experiment used: | Component | Version or pin | |---|---| | Python | `3.13.14` | | PyTensor | `3.1.2` | | PyTensor-ML | `f6ecf81d58da180cce50b77a43cf5d2c3d95e470` | | MLX | `0.32.0` | | Numba | `0.65.1` | | MLX-LM | `0.31.3` | | Machine | Apple M3 Max, 128 GB unified memory | The exact environment is recorded in [`environment.yml`](environment.yml). Code cells are not executed during the website build because the model artifacts are intentionally excluded from Git; the committed JSON reports preserve the artifact hashes and measured outputs. The PyTensor suite passed 934 tests with 30 skips, and the gated real-model runner tests passed all 189 checks; the compact record is in [`gemma3n_test_summary.json`](results/gemma3n_test_summary.json). --- The public package, model-specific implementations, tests, generated draft, audits, and sanitized result reports are available in the [site repository](https://github.com/cetagostini/cetagostini.github.io).