Mapped inputs, shared inputs and output axes

One vector for every row, or one offset per row?

Our batch has rows [1, 2, 3] and [4, 5, 6]. Adding a shared vector [10, 20, 30] gives [11, 22, 33] and [14, 25, 36]. Adding a separate scalar offset to each row, 100 then 200, gives [101, 102, 103] and [204, 205, 206]. Both operations are additions; the mapping specification says which values travel together.

The multi-input Vmap.apply accepts a function over a list of tensor arguments, an axis specification for each input, and one output axis. An integer input axis selects the dimension to map. A null input axis means the whole tensor is shared with every logical example.

Choose Run all, then alter axes and values in the individual cells. This notebook uses jax4j 0.2.1 on the emulated CPU. Each cell creates and closes its own ordinary tensor inputs and new outputs. Imports and the explicitly typed Java data array persist. Leaving or reloading discards state and edits.

1. Share a feature-wise vector

The batch is mapped on axis zero, while the bias vector is shared. Inside the logical function, the row and the shared bias both have shape [3]. Their addition is ordinary elementwise arithmetic.

Use Arrays.asList(0, null) for the input-axis list. Java's List.of rejects null entries before vmap even sees them; this small collection detail matters when spelling “shared” in Java.

Expect shape [2, 3] and flat values [11, 22, 33, 14, 25, 36]. Try this: change only the last bias value from 30 to 0. The third entry of both output rows changes. You are modifying one feature's shared parameter, so both observations should respond.

A shared input remains a tensor argument. It can carry new values on the next call; it is not a Java scalar frozen during JIT capture. If this mapped function is later compiled as a multi-leaf tree, the shared tensor is still an explicit runtime leaf in that outer function.

2. Map two arguments together

Now the second argument is a vector of two row offsets. Mapping its axis zero produces one scalar slice per example. That scalar broadcasts across the three entries of the corresponding logical row. Mapped axes must have the same length, but their remaining sample shapes can differ.

Expect [101, 102, 103, 204, 205, 206]. Change only the second offset to -100; only the second row changes. This is a different dependency pattern from the shared feature bias. Predicting which positions move is a useful check before worrying about the exact numbers.

The two mapped inputs are paired by index. This call combines row zero with offset zero and row one with offset one; it does not form every possible row-offset pair. To request a Cartesian product, you would need another mapping level or another explicit arrangement of axes.

3. Move the mapped output axis

The output-axis argument chooses where the mapping dimension appears in the returned tensor. The per-example addition still returns a vector of shape [3]. Putting the mapped dimension at output axis one produces shape [3, 2], with the two observations side by side for each feature.

Expect [3, 2] and [11, 14, 22, 25, 33, 36]. The output entries are rearranged; their numerical meaning has not changed. Compare this with the first cell's row-major output and draw the two matrices before experimenting with another axis value.

4. Map an input axis other than zero

We can map the original batch's axis one instead. There are now three logical examples, each containing two entries: the columns [1, 4], [2, 5], [3, 6]. A shared vector [10, 100] aligns with those two-entry examples. Putting the output mapped axis back at one preserves the original matrix layout.

Expect [2, 3] and [11, 12, 13, 104, 105, 106]. Try output axis zero instead. The result becomes [3, 2], with flat values [11, 104, 12, 105, 13, 106]. The logical examples are still columns; only their placement in the returned tensor changes.

Negative axes count from the end. Here input axis -1 means axis one of a rank-two batch, and output axis -1 also means axis one of its rank-two result. Try replacing both ones with -1. Normalize axes against the actual input or output rank, rather than treating a negative index as a special batching mode.

5. Return a mapped result and an unmapped summary

Vmap.applyTree supports structured outputs with one output-axis entry per leaf. Our function returns the corrected rows and the sum of the shared bias. That bias sum does not depend on a mapped input, so its output axis can be null. It is returned once, with shape [1], instead of being repeated per row.

Expect corrected shape [2, 3] and summary shape [1] with value 60. The summary is the sum of the bias, not the total of all corrected readings. Try replacing the second expression with a.get(1).mul(a.get(1)).sum(); the shared summary becomes 1400.

A null output axis requires an unmapped result. It cannot remove a genuine mapped dimension by declaration. Giving the corrected rows a null output axis would be rejected; reduce the mapped result explicitly afterward if you need a total across examples. Axis lists follow deterministic tree-leaf order, which matters when replacing this simple output list with records or maps.

6. Catch unequal mapped lengths

The following offsets have length three, while the batch's mapped row axis has length two. Because both inputs are marked mapped on axis zero, their lengths disagree before the body can combine them. Catching that error keeps Run all useful while showing the exact boundary being checked.

Keep the first input mapped and replace List.of(0, 0) with Arrays.asList(0, null). The three offsets become a shared feature vector, producing rows [101, 202, 303] and [104, 205, 306]. This fixes the shape mismatch by asking another question. Check that the new dependency pattern is actually what you intended.

Continue with explicit random keys before adding randomness to a mapped function. Revisit parameter trees when organizing several mapped and shared tensors into a readable input schema.