talks
|
texts
IRONWOOD
Copenhagen, Denmark
September 18, 2026
A short talk on annotation-driven vectorization for jaxtyped JAX dataclasses.
IRONWOOD
Noah Syrkis
September 22, 2026
1 |
Vectorizing structured JAX code
2 |
Annotations define vectorization
3 |
PyTrees and broadcasting
4 |
Controlling what “scalar” means
5 |
From one device to a mesh
6 |
The Ironwood programming model
1 |
Vectorizing structured JAX code
Consider an array function, mapping an
x
to
y
. Vectorizing across the leftmost dimension means
applying
f
independently to each row of
x
.
2 of
18
1 |
Vectorizing structured JAX code
import
jax
.
numpy
as
jnp
from
jax
import
Array
def
f
(
x
:
Array
)
->
Array
:
return
x
*
x
Listing 1: Libraries are important and a function
to be vectorized is defined
#
[0, 1, 2, ..., 9]
x
=
jnp
.
arange
(
10
)
#
[0, 1, 4, ..., 81]
y
=
jax
.
vmap
(
f
)
(
x
)
Listing 2: We declare an array for each entry of
which we want to call
f
.
3 of
18
1 |
Vectorizing structured JAX code
x[0]
x[1]
x[2]
f
y[0]
y[1]
y[2]
4 of
18
1 |
Vectorizing structured JAX code
Broadcasting shares an input across mapped
calls.
def
add
(
x
,
bias
)
:
return
x
+
bias
y
=
jax
.
vmap
(
add
,
in_axes
=
(
0
,
None
)
)
(
x
,
bias
)
x
varies.
bias
is reused.
x[i]
bias
add
y[i]
This is the first bookkeeping step: deciding which
leaves are mapped and which leaves are shared.
5 of
18
1 |
Vectorizing structured JAX code
jnp.vectorize
generalizes the idea.
It lets scalar signatures describe how arrays
should broadcast and map.
@
jnp
.
vectorize
(
signature
=
"
(d),(d)->()
"
)
def
dot
(
x
,
y
)
:
return
x
@
y
The signature says
d
is part of one scalar call.
Other leading dimensions broadcast and vector
ize.
x
.
shape
#
[batch, d]
y
.
shape
#
[d]
dot
(
x
,
y
)
.
shape
#
[batch]
6 of
18
1 |
Vectorizing structured JAX code
Ironwood applies that idea to PyTrees.
Real JAX code is often dataclass-shaped, not just
array-shaped.
@
dataclass
class
Example
:
x
:
Shaped
[
Array
,
"
*batch feature
"
]
mask
:
Shaped
[
Array
,
"
*batch
"
]
table
:
Shaped
[
Array
,
"
kind feature
"
]
The annotations say what
jnp.vectorize
signa
tures say for arrays.
They also work through structured JAX values.
7 of
18
2 |
Annotations define vectorization
Ironwood reads
jaxtyping
annotations.
x
:
Shaped
[
Array
,
"
*batch feature
"
]
y
:
Shaped
[
Array
,
"
*batch
"
]
z
:
Shaped
[
Array
,
"
feature
"
]
*batch
means: these leading axes are mapped.
[b, n]
*batch = b
feature = n
Everything after
*batch
belongs to the scalar
function.
8 of
18
2 |
Annotations define vectorization
The function stays ordinary.
from
ironwood
import
vectorize
@
vectorize
def
normalize
(
x
:
Shaped
[
Array
,
"
*batch
feature
"
]
)
:
return
x
/
jnp
.
linalg
.
norm
(
x
)
The scalar body sees one
feature
vector.
The caller may pass one vector, a batch of vec
tors, or a rectangular batch of vectors.
The output gets the same batch prefix back.
9 of
18
3 |
PyTrees and broadcasting
Dataclasses are PyTrees.
Ironwood descends through them, finds anno
tated leaves, and reconstructs the same structure
on return.
@
dataclass
class
Pair
:
left
:
Shaped
[
Array
,
"
*batch dim
"
]
right
:
Shaped
[
Array
,
"
*batch dim
"
]
@
vectorize
def
distance
(
pair
:
Pair
)
->
Array
:
return
jnp
.
linalg
.
norm
(
pair
.
left
-
pair
.
right
)
10 of
18
3 |
PyTrees and broadcasting
Leaves without a batch prefix broadcast.
@
dataclass
class
Batch
:
x
:
Shaped
[
Array
,
"
*batch dim
"
]
center
:
Shaped
[
Array
,
"
dim
"
]
x[b, d]
center[d]
scalar call
out[b]
One value can be shared by every mapped case
without spelling out an axes policy.
11 of
18
4 |
Controlling what “scalar” means
Scalar does not mean rank zero.
It means one outer case.
points
:
Shaped
[
Array
,
"
*batch point xy
"
]
point
and
xy
remain inside the scalar function.
@
vectorize
def
centroid
(
points
:
Shaped
[
Array
,
"
*batch point
xy
"
]
,
)
->
Shaped
[
Array
,
"
xy
"
]
:
return
points
.
mean
(
axis
=
0
)
The body computes one centroid. The wrapper
computes many.
12 of
18
4 |
Controlling what “scalar” means
Extend
keeps a named inner axis on the result.
from
ironwood
import
Extend
@
vectorize
def
classify
(
x
:
Shaped
[
Array
,
"
*batch item
feature
"
]
,
)
->
Extend
[
Shaped
[
Array
,
"
"
]
,
"
item
"
]
:
return
model
(
x
)
*batch
item
feature
scalar body
*batch item
The annotation says which axes are outer batch
and which axes are retained.
13 of
18
5 |
From one device to a mesh
The same annotated boundary can be sharded.
@
vectorize
(
shard_axis
=
"
data
"
)
def
score
(
x
:
Example
)
->
Result
:
return
scalar_score
(
x
)
batch
flatten
shard_map
vmap
restore
The scalar body is still local code.
14 of
18
5 |
From one device to a mesh
The mesh is explicit JAX state.
mesh
=
jax
.
make_mesh
(
(
4
,
)
,
(
"
data
"
,
)
)
with
jax
.
set_mesh
(
mesh
)
:
out
=
jax
.
jit
(
score
)
(
examples
)
Ironwood shards the mapped batch.
Non-batched leaves are replicated.
The batch size must divide cleanly across the
mesh axis.
15 of
18
6 |
The Ironwood programming model
Ironwood is a small adapter:
▶
read jaxtyping annotations
▶
flatten PyTrees
▶
broadcast scalar leaves
▶
call
vmap
▶
optionally call
shard_map
▶
rebuild the output
jaxtyped dataclass
Ironwood
vmap
shard_map
same shape
16 of
18
6 |
The Ironwood programming model
The rule:
Write the scalar program.
Put batch structure in the annotations.
Keep data structure in dataclasses.
Let JAX compile the transformed function.
@
vectorize
(
shard_axis
=
"
data
"
)
def
f
(
x
:
Example
,
key
:
Shaped
[
Array
,
"
*batch key
"
]
)
:
return
scalar_f
(
x
,
key
)
Ironwood is for code that is already pure, typed,
and structured.
17 of
18
References
