Vectorize with vmap

Use vmap to turn a function for one pair of vectors into a function for a batch of pairs. Each row of A is paired with the same row of B.

Your task

Implement batched_dot(A, B) using torch.vmap or jax.vmap, without an explicit Python loop. Both inputs are floating-point matrices of shape (batch, n) with the same dtype and device. Return shape (batch,), preserving dtype and device and leaving the inputs unchanged.

output[i] = sum(A[i, j] * B[i, j] over all j)

For A=[[1, 2], [3, 4]] and B=[[5, 6], [7, 8]], return [17, 53]: the first pair gives 1*5 + 2*6, and the second gives 3*7 + 4*8. Do not compute all pairs of rows or sum the batch into one scalar.

Define an inner function that takes two vectors of shape (n,) and returns a scalar dot product. Then call torch.vmap(dot)(A, B) or jax.vmap(dot)(A, B). The default mapping uses axis 0 of both inputs and places the mapped results on output axis 0. A batch of one still returns (1,).

Follow the shapes

For two (8, 3) inputs, the inner function sees two (3,) vectors. Each scalar result contributes one entry to the (8,) output. vmap transforms the tensor operations to handle the batch; it does not require you to call the Python function separately for every row.

The diagram starts with both inputs mapped over axis 0. [] denotes a scalar shape, not an empty array. Its controls explore axis placement; they do not check whether the inner dot product accepts every selected shape.