# PyMC Graph accepting gradients?

**URL:** <https://discourse.pymc.io/t/pymc-graph-accepting-gradients/14775>\
**Category:** v5\
**Created:** [July 15, 2024, 5:13pm UTC](https://discourse.pymc.io/t/pymc-graph-accepting-gradients/14775 "2024-07-15T17:13:50Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![AstroParrotKing](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/astroparrotking/32/7863_2.png) [@AstroParrotKing](https://discourse.pymc.io/u/AstroParrotKing)\
**Post date:** [July 15, 2024, 5:13pm UTC](https://discourse.pymc.io/t/pymc-graph-accepting-gradients/14775/1 "2024-07-15T17:13:50Z")

</div>

I am trying to understand how the wider PyMC graph accepts custom gradients from PyTensor Ops. I have my own PyTensor Op to calculate gradients, but I have realized that the wider PyMC model I have isn’t accepting them (or doesn’t recognize them).

In [this](https://www.pymc.io/projects/examples/en/latest/howto/blackbox_external_likelihood_numpy.html#pytensor-op-with-gradients) example can someone explain how exactly the gradients calculated get sent to the wider PyMC model? I assumed it was in this line:

```auto
return [
            pt.sum(out_grad * grad_wrt_m),
            pt.sum(out_grad * grad_wrt_c),
            # We did not implement gradients wrt to the last 3 inputs
            # This won't be a problem for sampling, as those are constants in our model
            pytensor.gradient.grad_not_implemented(self, 2, sigma),
            pytensor.gradient.grad_not_implemented(self, 3, x),
            pytensor.gradient.grad_not_implemented(self, 4, data),
        ]

```

since I do not understand this line:

```auto
[out_grad] = g

```

For added context my full Op is:

```auto
def jax_ecc_anom(manom, ecc):
    # manom = np.reshape(manom,[7])
    # print(type(ecc))
    alpha = (1.0 - ecc) / ((4.0 * ecc) + 0.5)
    beta = (0.5 * manom) / ((4.0 * ecc) + 0.5)

    aux = jnp.sqrt(beta **2.0 + alpha** 3.0) 
    z = jnp.abs(beta + aux)**(1.0/3.0) 

    s0 = z - (alpha/z)
    s1 = s0 - (0.078*(s0**5.0)) / (1.0 + ecc)
    e0 = manom + (ecc * (3.0*s1 - 4.0*(s1**3.0)))

    se0 = jnp.sin(e0)
    ce0 = jnp.cos(e0)

    f = e0-ecc*se0-manom

    f1 = 1.0-ecc*ce0
    f2 = ecc*se0
    f3 = ecc*ce0
    f4 = -f2
    u1 = -f/f1
    u2 = -f/(f1+0.5*f2*u1)
    u3 = -f/(f1+0.5*f2*u2+(1.0/6.0)*f3*u2*u2)
    u4 = -f/(f1+0.5*f2*u3+(1.0/6.0)*f3*u3*u3+(1.0/24.0)*f4*(u3**3.0))

    return (e0 + u4)

jitted_ecc = jax.jit(jax_ecc_anom)

def grad_ecc(EA, ecc):

    sea = jnp.sin(EA)
    cea = jnp.cos(EA)
    temp = 1 - ecc * cea
    dEA_dM = 1.0 / temp
    dEA_de = (sea * EA) / temp
    return dEA_dM, dEA_de

jitted_grad_ecc = jax.jit(grad_ecc)

class KeplerSolverOp(Op):
    __props__ = ()

    def make_node(self, M, e):
        M = at.as_tensor(M)
        e = at.as_tensor_variable(e)

        inputs = [M, e]
        outputs = [M.type()]

        return Apply(self, inputs, outputs)

    def perform(self, node, inputs, outputs):
        manom, ecc = inputs

        EA = jitted_ecc(manom, ecc) # Compute ecc-anomly

        outputs[0][0] = np.asarray(EA)

    def grad(self, inputs, out_grad):
        M, e = inputs

        grad_wrt_m, grad_wrt_e = kepler_loglikegrad(M,e)

        return[grad_wrt_m,grad_wrt_e]
    
class LogLikeGrad(Op):
    def make_node(self, M, e):
        M = at.as_tensor(M)
        e = at.as_tensor_variable(e)

        inputs = [M, e]

        outputs = [M.type(), e.type()]

        return Apply(self, inputs, outputs)

    def perform(self, node, inputs, outputs):
        manom, ecc = inputs

        EA = jitted_ecc(manom, ecc)

        dEA_dM, dEA_de = jitted_grad_ecc(EA,ecc)
        
        outputs[0][0] = dEA_dM
        outputs[1][0] = dEA_de

        # partials(M)/temp + partials(e)*sea/temp
        # dEA_dM, dEA_de, temp, sea = jitted_grad_ecc(EA,ecc)
        # outputs[0][0] = (dEA_dM/temp) + (dEA_de*sea)/(temp)

keplergrad = KeplerSolverOp()
kepler_loglikegrad = LogLikeGrad()

@jax_funcify.register(KeplerSolverOp) #done
def keplergrad_jaxify(op, **kwargs):#done
    return jax_ecc_anom#done

@jax_funcify.register(LogLikeGrad) #done
def kepler_loglikegrad_jaxify(op, **kwargs):#done
    return grad_ecc

```

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [July 15, 2024, 7:23pm UTC](https://discourse.pymc.io/t/pymc-graph-accepting-gradients/14775/2 "2024-07-15T19:23:57Z")

</div>

> [@AstroParrotKing](#):
>
> `[out_grad] = g`

That line is doing unpacking from a list with a single element.

For a single element you need to distinguish between the two cases

```python
x = [1] # x is [1]
[x] = [1] # x is 1
# same as (x,) = [1]
assert x == 1

```

You can also unpack multiple elements, which is commonly seen in loops, function calls

```python
[x, y] = [1, 2]
# same as x, y = [1, 2]
assert x + y == 3

```

Anyway, Python syntax aside, PyMC requests your gradient expressions when you (or PyMC) asks for the gradient of the logp. You can trigger this manually with `model.dlogp()`. Note that if you are using an external JAX sampler like `numpyro` or `blackjax` the gradient won’t be provided by PyTensor but instead by JAX, and it doesn’t matter whether you implemented it or not in PyTensor.

For a more intro picture of how PyMC gets to probabilites (and then its gradients): [PyMC and PyTensor — PyMC 5.16.2 documentation](https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/pymc_pytensor.html#enough-with-random-variables-i-want-to-see-some-log-probabilities)

---

<div class="post-metadata">

**Author:** ![AstroParrotKing](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/astroparrotking/32/7863_2.png) [@AstroParrotKing](https://discourse.pymc.io/u/AstroParrotKing)\
**Post date:** [July 16, 2024, 5:57pm UTC](https://discourse.pymc.io/t/pymc-graph-accepting-gradients/14775/3 "2024-07-16T17:57:23Z")

</div>

So my goal in creating the previous PyTensor Op was to feed the analytical gradient made from `jitted_grad_ecc` into PyMC’s wider graph as a sort of “differentiation rule” for only this specific this `jitted_ecc` function while sampling from `numpyro`

Because I noticed that that wasn’t happening, would adding something like

```auto
logp_fn = model.compile_fn(model.logp(sum=False), mode='JAX')

```

after the model is stated work?

I have checked out [How to use JAX ODEs and Neural Networks in PyMC - PyMC Labs](https://www.pymc-labs.com/blog-posts/jax-functions-in-pymc-3-quick-examples/#:~:text=Wrapping%20a%20pure%20JAX%20function) and te link you sent before and I think that should work? Would I be missing anything?

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [July 17, 2024, 5:33am UTC](https://discourse.pymc.io/t/pymc-graph-accepting-gradients/14775/4 "2024-07-17T05:33:56Z")

</div>

> [@AstroParrotKing](#):
>
> while sampling from `numpyro`

PyMC/PyTensor are not involved in the gradients when using numpyro. If you want to overload the gradients numpyro uses you’ll have to do it with jax custom gradient API: [Custom derivative rules for JAX-transformable Python functions — JAX documentation](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html)

If you notice the last example in the blogpost we don’t bother defining gradients in PyTensor because we’ll sample with numpyro