Invariant of gradient operation broken?

[scroyston]

Still working on a smaller test case I can post, but shouldn't all 3 check_grads() pass independent of pred_func definition?
First 2 pass, 3rd fails.
I believe it fails w.r.t 'scale', as when I calc that grad directly I get a very different answer than grad(full_loss, (0,1))(...)
Any insight on how to debug would be greatly appreciated, thanks!:

#params = {'a': 1.67, 'b': 8.26, etc...}
#scale = 1.23e-01
#tm = {'input': big NxM, 'tgt': Nx1}
#pred_func, arbitrary function NxM -> Nx1

def full_loss(pms, scl):
    X = pred_func(pms, tm)
    sX = X * scl
    return jnp.mean((sX - tm['tgt']) ** 2)

def loss_test1(pms):
    return full_loss(pms, scale)

def loss_test2(scl):
    return full_loss(params, scl)

from jax.test_util import check_grads
check_grads(loss_test1, (params,), order=1)
print("pass 1")
check_grads(loss_test2, (scale,), order=1)
print("pass 2")
check_grads(full_loss, (params, scale), order=1)

[mattjj]

Thanks for raising this!

It may be hard to debug without a runnable repro... Can you share some numerical values, like how different 'very different' is? I'm wondering if the failure is just numerical. (After all, grad(full_loss)(...) and grad(full_loss, (0, 1))(...) are calculated almost identically...)