Unexpected Output with jax.grad

[ijku]

Greeting community,
I'm having some issues when using jax.grad. For simplicity, the term I'd like to compute: $\frac{\partial F_i}{\partial\eta_i}$ where $F_i=\frac{\eta_i^2 \rho_i}{\sum_{i=1}^{N}\eta_i^2}$
If we perform the derivation manually, the result should be:
$\frac{\partial F_i}{\eta_i}=\frac{2\eta_i\sum_{i=1}^{N}\eta_i^2-(2\eta_i)\eta_i^2\rho_i}{(\sum_{i=1}^{N}\eta_i^2)^2}=\frac{2\eta_i}{\sum_{i=1}^{N}\eta_i^2}(\rho_i-F_i)$

Below is the code I wrote, but the output isn't correct:

import jax
import jax.numpy as jnp

def f(eta, rho):
    return jnp.sum( (eta**2 * rho)/jnp.sum(eta**2) )

grad_f = jax.grad(f, argnums=0)

# test data
eta = jnp.array([0.2, 0.8, 0])
rho = jnp.array([1, 10, 0])

print(grad_f(eta, rho))
# this output is: [-4.9826984  1.2456741  0.       ]

The correct output can be obtained using the following code:

import jax
import jax.numpy as jnp

def f(eta, rho):
    return (eta**2 * rho) / jnp.sum(eta**2)

jac_f = jax.jacobian(f, argnums=0)

# test data
eta = jnp.array([0.2, 0.8, 0])
rho = jnp.array([1, 10, 0])

jac_f_val = jac_f(eta, rho)
result = jnp.diag(jac_f_val)
print(result)
# this output is: [0.5536332 1.3840828 0.       ], which is correct

The computed result [0.5536332 1.3840828 0. ] matches the manually derived result previously mentioned.

# manual
print(2*0.2 / (0.2**2+0.8**2) * (1-(0.2**2*1)/(0.2**2+0.8**2)))  # 0.5536332179930795
print(2*0.8 / (0.2**2+0.8**2) * (10-(0.8**2*10)/(0.2**2+0.8**2)))  # 1.384083044982698

1. Could someone please help me identify what's wrong with my initial code?
2. is there any trick to ensuring that jax.grad computes the correct result?

[jakevdp]

1. Could someone please help me identify what's wrong with my initial code?

I think both results are correct, in that JAX is computing the requested gradient for the function you defined.

Note that if you change your second output to this:

result = jnp.sum(jac_f_val, axis=0)

then you get the same result as the first. This makes sense because summation and differentiation commute: the gradient of the sum will be the same as the sum of the gradients (but will not be the same as a single term from the sum).

2. is there any trick to ensuring that jax.grad computes the correct result?

I'm not sure there's any trick beyond making sure you're writing code which matches the mathematical expression you have in mind!

Hope that helps!

[ijku]

Thank you for your response!

For the 1st point, I found that it actually computes the gradient of $\frac{\partial}{\partial \eta_i}(\frac{\sum_{i=1}^{N}\eta_i^2 \rho_i}{\sum_{i=1}^{N}\eta_i^2})$ —that is, it effectively introduces a summation symbol in the numerator of the original expression. The partial derivative is:

$\frac{\partial}{\partial \eta_i}(\frac{\sum_{i=1}^{N}\eta_i^2 \rho_i}{\sum_{i=1}^{N}\eta_i^2})=\frac{2\eta_i\rho_i\sum_{i=1}^{N}\eta_i^2-(2\eta_i)\sum_{i=1}^{N}\eta_i^2\rho_i}{(\sum_{i=1}^{N}\eta_i^2)^2}=\frac{2\eta_i}{\sum_{i=1}^{N}\eta_i^2}(\rho_i-\frac{\sum_{i=1}^{N}\eta_i^2\rho_i}{\sum_{i=1}^{N}\eta_i^2})$

In other words, the np.sum() in my initial function's return, which makes the result a scalar, ultimately changed the way my function was defined.

Here is the corresponding code:

import jax
import jax.numpy as jnp

def f(eta, rho):
    return jnp.sum((eta**2 * rho))/jnp.sum(eta**2) 

grad_f = jax.grad(f, argnums=0)

# test data
eta = jnp.array([0.2, 0.8, 0])
rho = jnp.array([1, 10, 0])

print(grad_f(eta, rho))
# output: [-4.9826984  1.2456741  0.       ]

# manually
print(2*0.2 / (0.2**2+0.8**2) * (1-(0.2**2*1+0.8**2*10)/(0.2**2+0.8**2) ))  # -4.982698961937715
print(2*0.8 / (0.2**2+0.8**2) * (10-(0.2**2*1+0.8**2*10)/(0.2**2+0.8**2) ))  # 1.2456747404844286

[jakevdp]

In other words, the np.sum() in my initial function's return, which makes the result a scalar, ultimately changed the way my function was defined.

Yes - that makes sense!