A dot product combines two same-length vectors by multiplying corresponding elements and summing the results. In JAX, jnp.vdot handles the operation directly. section 6.1
A small JAX calculation
def dot(v1, v2): return jnp.vdot(v1, v2) dot(jnp.array([1., 1., 1.]), jnp.array([1., 2., -1.])) >>> Array(2., dtype=float32) The result is 2: 1×1 + 1×2 + 1×−1. section 6.1
A function written for one vector pair does not automatically accept a batch. Passing 32 images can produce a contracting-dimensions error because the batch dimension is mistaken for part of the dot product. section 2.5
Listing 6.1 A function for calculating a dot product between two vectors
def dot(v1, v2): return jnp.vdot(v1, v2) #1 dot(jnp.array([1., 1., 1.]), jnp.array([1., 2., -1])) #2 >>> Array(2., dtype=float32)
Inspect the compact JAX implementation and its numerical result. section 6.1
