Skip to content

← JAX Profiling & Debugging step 7 of 12

Medium Primitives

Append a Scalar with Concatenate

Jump to the task

Append a scalar without unpacking the array

To append one value to a vector, concatenate the vector with a one-element array. This expresses the operation directly and keeps the input unchanged.

import jax
import jax.numpy as jnp

x = jnp.array([1., 2., 3.])
val = jnp.array(4.)              # shape ()
tail = jnp.atleast_1d(val)       # shape (1,)
out = jnp.concatenate([x, tail])
# [1., 2., 3., 4.], shape (4,)

Concatenate joins along an existing axis, here axis 0. Both inputs must have matching rank. A scalar has rank zero, so promote it to a one-element vector first. stack would add an axis and is not the operation requested here.

Local lists are not automatically tracer leaks

This alternative can work under ordinary JIT when x has a static leading dimension:

@jax.jit
def list_version(x, val):
    pieces = list(x)
    pieces.append(val)
    return jnp.array(pieces)

Python builds the local list while tracing. Its elements can be tracers that still represent runtime values; appending one does not itself request a boolean or freeze the first input's numbers. Subsequent cached calls can produce different values correctly.

The drawback is that unpacking a large array creates per-element Python work and a more elaborate traced program. Concatenation avoids that unpacking. A real leak occurs if traced values escape through outside state, such as a global list, rather than normal return values.

jnp.append also supports this numerical operation. There is no universal guarantee that it is slower or uses an extra physical allocation. This exercise specifically practices concatenate and rank handling.

Problem

Implement append_via_concat(x, val) using jnp.concatenate.

X is a one-dimensional float32 JAX array. Val is a Python float or a scalar float32 JAX array. Return a float32 JAX array of shape (N + 1,), where N is x.shape[0], preserving x's entries and placing val last. An empty x produces a one-element array containing val. Do not mutate x.

The function must work under jax.jit with dynamic val. Its output size depends on the static input shape, not on val's numerical value. Changing the input length can cause retracing. Output tests check the values and shape; inspect that you used the requested concatenation.