I am unsure how to evaluate the vector-Jacobian-product using non-identity vectors in JAX.

edit

To evaluate the vector-Jacobian product using non-identity vectors in JAX, you can use the jax.vjp function. This function takes a function and an input value, and returns two values: the output of the function and a function that computes the vector-Jacobian product.

Here's an example:

import jax.numpy as np
from jax import grad, jit, vmap, vjp

# Define a function
def f(x):
  return np.array([x[0] * x[1], np.sin(x[0])])

# Evaluate the function and its Jacobian at an input value
x = np.array([1.0, 2.0])
y, jacobian = vjp(f, x)

# Define a non-identity vector to compute the vector-Jacobian product
v = np.array([1.0, 2.0])

# Compute the vector-Jacobian product
vjp_product = jacobian(v)

print("Function output:", y)
print("Jacobian:\n", jacobian(np.eye(2)))
print("Vector:", v)
print("VJP product:", vjp_product)

In this example, f() is a function that takes a 2-dimensional input and returns a 2-dimensional output. We use vjp() to compute the function output and its Jacobian at the point x = [1.0, 2.0]. Then, we define a non-identity vector v = [1.0, 2.0] and compute the vector-Jacobian product using the jacobian() function returned by vjp(). Finally, we print the function output, the Jacobian, the vector, and the vector-Jacobian product.

Note that in the above example, we define v as a 2-dimensional vector. If f had a higher-dimensional output or input, v would need to match the dimensions of the relevant vector/derivative.