Skip to content

← JAX Foundations step 19 of 25

Easy Primitives

Translate a PyTorch Matrix Pipeline to JAX

Jump to the task

Why this matters

JAX and PyTorch offer many similar array operations, but translating code means checking shapes, axis arguments and behavior, not just changing a namespace. This exercise translates a short matrix pipeline and preserves the order of its output values.

PyTorch operation JAX counterpart
torch.matmul(a, b) jnp.matmul(a, b) or a @ b
torch.where(cond, t, f) jnp.where(cond, t, f)
torch.zeros_like(t) jnp.zeros_like(t)
torch.transpose(t, i, j) jnp.swapaxes(t, i, j)
t.permute(2, 0, 1) jnp.transpose(t, (2, 0, 1))
t.unsqueeze(0) jnp.expand_dims(t, 0)
torch.cat(parts, dim=0) jnp.concatenate(parts, axis=0)

These are operation mappings, not promises of identical dtype promotion, memory aliasing, device placement or differentiation behavior in every case.

Worked mini-example

import jax.numpy as jnp

x = jnp.array([[1.0, 0.0], [0.0, 1.0]])
w = jnp.array([[2.0, -1.0], [3.0, 4.0]])
y = x @ w                              # [[2, -1], [3, 4]]
y = jnp.where(y > 0, y, jnp.zeros_like(y))  # [[2, 0], [3, 4]]
y = jnp.transpose(y, (1, 0))            # [[2, 3], [0, 4]]
result = y.reshape(-1)                  # [2, 3, 0, 4]

Flattening before transposing would instead give [2, 0, 3, 4]. The transpose changes which values become neighbors in the output.

Problem

Implement torch_to_jax_translation(x, w), the JAX equivalent of:

def torch_pipeline(x, w):
    import torch
    y = torch.matmul(x, w)
    y = torch.where(y > 0, y, torch.zeros_like(y))
    y = torch.transpose(y, 0, 1)
    return y.reshape(-1)

x and w are finite float32 matrices of shapes (N, d_in) and (d_in, d_out). Assume intermediate arithmetic stays finite. Return a float32 vector of shape (N * d_out,): matrix multiply, replace nonpositive entries with zero, transpose the result's two axes, then flatten in row-major order.

Common pitfalls

  • torch.transpose swaps two axes. Its direct JAX counterpart is jnp.swapaxes. jnp.transpose instead takes a complete permutation when axes are supplied. For a matrix, (1, 0) works; for a rank-three array, swapping axes 0 and 1 requires (1, 0, 2).
  • Matrix .T swaps its two axes. On a higher-rank JAX array it reverses all axes, so it is not a general two-axis swap.
  • JAX updates return new array values. Assign the result of an .at update; it does not mutate the original array like a PyTorch method whose name ends in an underscore.
  • Similar names can hide different edge cases. Squeezing a specified non-singleton axis raises in JAX, while PyTorch leaves it unchanged.
  • Random generation and autodiff need separate consideration. JAX uses explicit keys and function transformations; a PyTorch generator or no_grad context is not translated by a namespace replacement.