Reshape, slice and reduce

Keep the experiment's structure

You have two runs of an experiment. Each run contains three time steps, and each time step contains two channels. Storing twelve numbers is easy. Keeping track of what each number means while selecting and summarizing them is the interesting part.

We will use shape [2, 3, 2], with axes named run, time, channel in that order. Run zero contains pairs [1, 10], [2, 20], [3, 30]; run one contains [4, 40], [5, 50], [6, 60]. The second channel is deliberately ten times the first, making an accidental rearrangement easy to spot. These are constructed values, not recorded observations.

The previous lesson used broadcasting to match different shapes. Here, we choose a piece of a tensor, rearrange its axes, and reduce the axes we no longer need. Those are three different decisions, and small examples make the distinction visible.

Choose Run all, then revisit the cells with your own changes. Setup imports and the observations Java array persist. Each cell creates and closes its own tensor handles. The notebook uses the bundled jax4j 0.2.1 CPU backend; the browser runs an emulated JVM, so focus on results and shapes rather than its timings.

1. Name every axis before reducing it

The first dimension answers “which run?”, the second “which time?”, and the third “which channel?”. Size is 2 × 3 × 2 = 12; rank is three. The flat array groups the final axis most tightly, so each channel pair stays adjacent.

Print batch.toFloatArray() inside the block if you want to see the sequence again. The shape provides the grouping; the downloaded array is flat. Imagine placing a separator after every two entries and another after every three pairs. That is enough to reconstruct the two runs by hand.

Try this: change the last pair to [9, 90], rerun setup, and predict the mean of run one's channels before reaching the reduction cells. Restore [6, 60] to follow the reference answers. Later cells reuse the current Java array, not a hidden copy of the original example.

2. Slice the leading axis

sliceAxis0(1) selects the second run. Indices begin at zero. Selecting one position removes that leading axis, leaving shape [3, 2]: time and channel. This method selects one leading-axis entry; it is not Python-style range-slice syntax.

Expect the selected values [4, 40, 5, 50, 6, 60] and channel means [5, 50]. After slicing, axis zero refers to time. It no longer refers to run because that axis was removed. A surprisingly common mistake is carrying old axis numbers into a new shape without checking what remains.

Change the slice index to zero. The channel means become [2, 20]. Next, try changing run.sum(0) to run.sum(1): you now get one sum per time step instead of one sum per channel. Dividing those values by three would no longer compute the desired channel means. An expression can execute successfully while its interpretation has become wrong.

3. Reshape groups values; permute reorders axes

To combine runs and time into six rows, reshape(6, 2) is enough: the channel pairs already occur in the right order. To put time before run, use permute(1, 0, 2). Its arguments list the original axes in their new order. The result has shape [3, 2, 2].

The permuted flat sequence is [1, 10, 4, 40, 2, 20, 5, 50, 3, 30, 6, 60]. Each time step now groups both runs before moving to the next time. A plain reshape(3, 2, 2) would keep the old flat sequence and assign different meanings to neighboring values. Same dimensions, different data arrangement!

Try both operations and print their arrays side by side. Then use permute(0, 2, 1) to arrange run, channel, time. Which three values should appear first? They are the first channel's [1, 2, 3] within run zero. The method requires each original axis exactly once; duplicate axes cannot describe a permutation.

4. Reduce time and retain the channel distinction

We can compute both runs' channel means directly. Reduce axis one, the time axis of the original shape, and divide by three. Keeping the dimension gives shape [2, 1, 2]. That singleton time axis lets the mean broadcast back across the three time steps when we subtract it.

The means are [2, 20, 5, 50]. Each centered run becomes [-1, -10, 0, 0, 1, 10]. Sum a centered run along time and each channel should be zero for these exact integer inputs. You have combined a reduction with a broadcast to normalize a structured batch, without a Java loop over runs.

Try changing true to false. The means then have shape [2, 2], which does not broadcast to [2, 3, 2]: the aligned time dimensions disagree. The notebook will show an exception. Restore true and rerun. Keeping a reduced axis is a way to preserve alignment, not a way to retain its original data.

You can also reduce every element with mean(). Here that gives 19.25, but it mixes channels with different scales and both runs together. It is a valid arithmetic result and a poor substitute for the per-channel summary we asked for. Decide which distinctions matter before reducing them away.

5. Maximum value or maximum location?

max(axis) returns values. argmax(axis) returns positions along that axis. For our time sequences, each channel reaches its maximum at time index two. The index refers to the selected axis, not to a position in the flattened Java array.

Expect peak values [3, 30, 6, 60] and indices [2, 2, 2, 2], both arranged as [run, channel]. In this version, argmax indices are represented as F32 values, so the Java printout includes decimal points. Small integer indices are exact in F32; the result is not an integer-typed tensor just because it represents positions.

Change run zero's first pair to [8, 80] and rerun setup and this cell. Its first two peak indices should move to zero; the last two, belonging to the unchanged second run, remain two. Print the peak values alongside the indices and check both. Axis labels make a better guide than the flat printout.

6. Keep selection mistakes recoverable

There are only two runs, so valid leading-axis indices are zero and one. The next cell catches an out-of-range selection. Because the failure is contained, Run all still completes and you can immediately try a valid index.

Change the index to one and the cell prints [3, 2]. If you want to select another axis, first decide whether reordering axes is appropriate; do not assume that the index argument also specifies which axis to slice. The current notebook examples deliberately use this narrow API so the shape change stays explicit.

For a final experiment, keep only run zero, reshape its values into a vector of length six, and compare its overall mean 11.0 with its two channel means [2, 20]. Both answers are correct for their respective questions. Now explain which operations changed order, which changed grouping, and which discarded information. If you can follow those changes, a twelve-number tensor has done its teaching job.

Continue with matrix multiplication and linear algebra to combine feature values through weighted sums. For alignment practice, return to broadcasting and give the runs a shared calibration vector before reducing time.