# Can jax be used to wrap a function that contained python libraries?

**URL:** <https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489>\
**Category:** v5\
**Tags:** jax\
**Created:** [December 19, 2023, 9:13am UTC](https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489 "2023-12-19T09:13:27Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![liu\_liu](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/liu_liu/32/7513_2.png) [@liu\_liu](https://discourse.pymc.io/u/liu_liu)\
**Post date:** [December 19, 2023, 9:13am UTC](https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489/1 "2023-12-19T09:13:27Z")

</div>

Hi all, I am trying to wrap my custom likelihood function using jax.  
I followed [How to use JAX ODEs and Neural Networks in PyMC](https://www.pymc-labs.com/blog-posts/jax-functions-in-pymc-3-quick-examples/) guidance.  
My custom likelihood function depends on the **cantera** module.

Part of my code are:

> def func1(theta):
> 
> > ```
> > sol = func(theta)
> > return sol
> > 
> > ```
> 
> jitted\_custom\_op\_jax = jax.jit(func1)

> def vjp\_custom\_op\_jax(x, gz):
> 
> > ```
> > _, vjp_fn = jax.vjp(func1, x)
> > return vjp_fn(gz)[0]
> > 
> > ```
> 
> jitted\_vjp\_custom\_op\_jax = jax.jit(vjp\_custom\_op\_jax)

When I run my code.  
I got:

> File “C:\Users\86173\Desktop\jax\_sofc.py”, line 651, in func1  
> **sol = func(theta)**  
> ^^^^^^^^^^^  
> File “C:\Users\86173\Desktop\jax\_sofc.py”, line 552, in func  
> spec.thermo = ct.ConstantCp(T\_low=spec.thermo.min\_temp, T\_high=spec.thermo.max\_temp,  
> P\_ref=spec.thermo.reference\_pressure,  
> coeffs=[298.15, **-params[2]**, 139600, 0.0])  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “build\python\cantera\speciesthermo.pyx”, line 42, in cantera.speciesthermo.SpeciesThermo. **cinit**  
> jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method **array** () was called on traced array with shape float64.  
> The error occurred while tracing the function func1 at C:\Users\86173\Desktop\jax\_sofc.py:650 for jit. **This concrete value was not available in Python because it depends on the value of the argument theta.**

**-params[2]** is a value of theta I passed to function.

I wonder if there is any limitation for wrapping jax function? Such as we can just use some specific modules like numpy, so I cannot user **cantera**?  
But in [How to use JAX ODEs and Neural Networks in PyMC](https://www.pymc-labs.com/blog-posts/jax-functions-in-pymc-3-quick-examples/) it is able to use diffrax.diffeqsolve. I am confused.

---

<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:** [December 20, 2023, 11:29am UTC](https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489/2 "2023-12-20T11:29:43Z")

</div>

> [@liu\_liu](#):
>
> diffrax.diffeqsolve

This is a JAX module with JAX functions.

You cannot use JAX to wrap numpy functions. You can however use PyTensor. In that notebook PyTensor was wrapping JAX functions (not the other way around) but it can actually wrap any Python functions.

The difference is that you will have to implement the gradients manually instead of relying on JAX to extract them. There’s an example here: [Using a “black box” likelihood function (numpy) — PyMC example gallery](https://www.pymc.io/projects/examples/en/latest/howto/blackbox_external_likelihood_numpy.html)

---

<div class="post-metadata">

**Author:** ![liu\_liu](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/liu_liu/32/7513_2.png) [@liu\_liu](https://discourse.pymc.io/u/liu_liu)\
**Post date:** [December 25, 2023, 9:25am UTC](https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489/3 "2023-12-25T09:25:57Z")

</div>

Got it. Thanks!