Gradient only defined for scalar-output functions. Output had shape: (1,).

[bengladwyn]

I'm trying to calculate the gradient of a function V(h) which is defined via V(phi) and h(phi). I therefore need to invert h(phi) to find phi(h) which I do with a minimize function. Then, to find V(h) I calculate V(phi(h)). This works to successfully plot the function but I get the following error when trying to calculate the gradient:

in minimize_bfgs
    f_0, g_0 = jax.value_and_grad(fun)(x0)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: Gradient only defined for scalar-output functions. Output had shape: (1,).

My code is as follows

import matplotlib.pyplot as plt
from jax import grad
from jax.scipy import optimize
import jax.numpy as np

# Define specific parameter values
xi = 0.0218
c2 = 0.03
c3 = 0.075
phi0 = 0.7
lam = 9e-10

# Define the potential V(phi)
def V_phi(phi):
    return lam / (24 * (1 + xi * phi**2)**2) * (
    3 * phi**4 + xi**2 * phi0**4 * phi**4 - 8 * (1 + c3) * phi0 * phi**3 +
    2 * (1 + c2) * (3 + xi * phi0**2) * phi0**2 * phi**2)

# Define the Weyl transformation and canonical normalization
def hofphi(phi):
    return np.sqrt((1 + 6 * xi) / xi) * np.arcsinh(phi * np.sqrt(xi * (1 + 6 * xi))) - \
           np.sqrt(6) * np.arctanh((np.sqrt(6) * xi * phi) / np.sqrt(1 + xi * (1 + 6 * xi) * phi**2))

# Function to find the root of
def root_function(phi, h):
    return hofphi(phi) - h

# Function to find phi for an arbitrary value of h
def find_phi_for_h(h):
    phi_initial_guess = np.array([1.0])
    solution = optimize.minimize(root_function, phi_initial_guess, args=(h,),method="BFGS")
    if solution.success:
        return solution.x[0]
    else:
        raise ValueError("Root finding did not converge.")

# Define the potential V(h) in terms of h
def V_h(h):
    phi = find_phi_for_h(h)
    return V_phi(phi)

# Generate values for h between 0 and 2
h_values = np.linspace(0, 2, 400)
V_h_values = [V_phi(h) for h in h_values]

# Plot the potential V(h)

plt.plot(h_values, V_h_values, label='V(h)')
plt.xlabel('h')
plt.ylabel('V(h)')
plt.title('Potential V(h) vs h')
plt.legend()
plt.show()

# Compute the gradient of V(h) with respect to h
gradient_V_h = grad(V_h)
# Generate gradient values for h between 0 and 2
V_h_gradient_values = [gradient_V_h(h) for h in h_values]

# Plot the gradient of the potential V(h)
plt.plot(h_values, V_h_gradient_values, label="V'(h)")
plt.xlabel('h')
plt.ylabel("V'(h)")
plt.title("Gradient of Potential V(h) vs h")
plt.legend()
plt.show()

[jakevdp]

Thanks for the question! The issue is that your function returns a one-dimensional vector, and grad is only compatible with functions that return a scalar (i.e. an array with zero-dimensions). You could make your root function return a scalar by changing it to this:

def root_function(phi, h):
    return hofphi(phi)[0] - h

Also, side note: in both cases where you use list comprehensions, it will generally be more efficient to use vmap instead:

from jax import vmap

# V_h_values = [V_phi(h) for h in h_values]
V_h_values = vmap(V_phi)(h_values)

# V_h_gradient_values = [gradient_V_h(h) for h in h_values]
V_h_gradient_values = vmap(gradient_V_h)(h_values)