lax function greater than return zero gradient ?

[organic-chemistry]

Hello,
I would like to write a program that draw from a source which is a 1 dimentionnal array with weights.
I compare these weights to a uniform distribution, and it the weights are higher than the uniform distribution I
draw a sample.
My goal is to fit the drawn sample to a target.
However the gradient is zero.

import jax.numpy as jnp
from jax import grad
from jax import lax
from jax import random

def update_source(source,key):
    r_arr = random.uniform(key,shape=[len(source)])
    new  = lax.gt(0.5* source, r_arr)
    return new


source = jnp.array([0,0,1,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,0,1,0],dtype=float)
target= jnp.array([0,0,1,0,1,0,1,0,0,3,0,0,0,0,1,0,1,0,0,0,0,1.])


def loss(oris):
    sample= update_source(source,key=random.PRNGKey(0))
    return jnp.mean( (sample-target)**2)

def gloss(oris):
    return grad(loss)(oris) 
print(gloss(source))

and it return

[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]

Is it possible to make this code differentiable?
Best

[jakevdp]

Zero is the correct gradient here.

Remember that what differentiation is measuring the change in f given an infinitesimal change in x. Let's write a slightly simpler version of the same kind of function:

from jax import grad

def f(x):
  return (x > 0).astype(float)

grad(f)(1.0)
# 0.0

For x=1.0, f returns 1.0. If we perturb x to make it slightly larger or smaller, this does not change the output, so by definition, grad(f)(1.0) should be zero. This same logic holds for all values of f greater than zero: infinitessimally perturbing the input does not change the output, so the gradient is zero.

Similarly, for all values of x less than zero, the output is zero. Perturbing x does not change this output, so the gradient is zero.

That leaves us with the tricky case of x=0. Surely, if you perturb x upward, it will change the output, but this is problematic: an infinitesimal change in x produces a finite change inthe function value, which implies the gradient is undefined. Fortunately, there's another way for us to measure the gradient in this case: we perturb the function downward, in which case the output does not change, and so the gradient is zero.

JAX and other autodiff systems tend to handle discontinuities in this way: if the positive gradient and negative gradient disagree, but one is defined and the other is not, we use the one that is defined.

Under this definition of the gradient, mathematically and numerically the gradient of this function is everywhere zero.

[organic-chemistry]

Ok, thank you for the clarity of your response. It helped me reformulate the problem.
Maybe I could use as a proxy for the gradient, the gradient of the average value, because for my final problem
what I fit is the average realisation of the function.
if I define this function:

from numpy import random
def f(x): 
    r_arr = random.uniform()
    if x> r_arr:
        return 1
    else:
        return 0

Then f(0.5) will return 0 or 1 but its average will be 0.5.
So that the gradient of the average is 1 between [0,1] and 0 elsewhere.

in jax code I managed to write this:

from jax import custom_jvp
import jax.numpy as jnp
from jax import grad
from jax import lax
from jax import random,vmap

@custom_jvp
def compare(v,r):
    return lax.gt(v,r)*1.0

@compare.defjvp
def compare_jvp(primals, tangents):
  x, y = primals
  x_dot, y_dot = tangents
  primal_out = compare(x, y)
  tangent_out = x_dot
  return primal_out, tangent_out

key=random.PRNGKey(0)
r_arr = random.uniform(key)

print(grad(compare)(0.1,r_arr)) # return 1

(I will take care of the values lower than 0 or higher than one in the boundary of my optimisation)

However it does not work if I try to vmap the compare function:

print(grad(vmap(compare))(jnp.array([0.1])[:],jnp.array([r_arr])[:]))
#return: TypeError: Gradient only defined for scalar-output functions. Output had shape: (1,).

[organic-chemistry]

Oh sorry I realised my mistake for the last part.
I should take

vmap(grad(compare))

[jakevdp]

That certainly does return a nonzero gradient, but I don't think it's the correct nonzero gradient for your function.

Your compare function is essentially a Heaviside step function; if you want a differentiable approximation to this, you could replace it with a sigmoid. This has the same behavior for large positive and negative inputs, but replaces the discontinuity with a smooth, differentiable curve.