Custom VJP for Dirac delta function

[pablo2909]

Hello,

I am working on some NeuralODE code and I am facing an issue that, I hope, can be understood from the minimal example below.

Assume I have the following function f and I want to define a custom_vjp rule:

@jax.custom_vjp
def f(p, x):
    """
    Args:
        p: A scalar parameter.
        x: A vector.
    """
    return p*x**2

def f_fwd(p,x):
    return f(p,x), x

def f_bwd(res, g):
    print(f"{g = }")
    print(f"{res = }")
    return jnp.multiply(res**2, g+1), None

f.defvjp(f_fwd, f_bwd)

I compute some loss on it as follow:

def L(p,x, s):
    """
    Args:
        p: A scalar parameter.
        x: A vector.
        s: A scalar start time.
    """
    return jnp.sum((f(p,x)[s:]-5))**2

So in the loss function, I am only interested about the elements indexed from s, and because of the slicing [s:], the cotangent g will be of the form [0,...,0,x,...,y] where the values are 0 until the index s. Now notice that in the return statement of f_bwd I am adding 1 to the vector g. Actually what I want to do is add 1 only for the element that are indexed from s. How can I detect in f_bwd the index s ?

Thank you very much for any help !

PS: The title is maybe a bit misleading, if anyone has a better proposition I would be happy to change !
PPS: To give a bit more background on my Neural ODE problem. The function f represents an ODE solver and I want to compute the gradients of the parameters using the adjoint method. The thing is that the solver outputs T time steps so the cotangent vector in the backward pass will be of length T. But I am interested only by the last s time steps of the solution to compute my loss, hence how do I detect in the cotangent those last s values.