The jaxfit from dipolar-quantum-gases

jaxfit's People

Contributors

Stargazers

Watchers

jaxfit's Issues

Using curve_fit with function that was already jitted

For example following code (based on example from jaxfit docs but with @jax.jit added)

import jax
from jaxfit import CurveFit

@jax.jit
def linear(x, m, b): # fit function
     return m * x + b

x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
y = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

cf = CurveFit()
popt, pcov = cf.curve_fit(linear, x, y)

fails with an error like

Traceback (most recent call last):
  File "/home/annulen/./jax_fit.py", line 14, in <module>
    popt, pcov = cf.curve_fit(linear, x, y)
  File "/home/annulen/venv/jax/lib/python3.10/site-packages/jaxfit/minpack.py", line 228, in curve_fit
    res = self.ls.least_squares(f, p0, jac=jac, xdata=xdata, ydata=ydata, 
  File "/home/annulen/venv/jax/lib/python3.10/site-packages/jaxfit/least_squares.py", line 469, in least_squares
    func_update = self.f.__code__.co_code != fun.__code__.co_code
AttributeError: 'PjitFunction' object has no attribute '__code__'. Did you mean: '__call__'?

However, built-in JAX functions like jax.vmap or jax.grad work fine with both compiled and non-compiled functions. It would be great if it worked here too.

ploynomial curve fiting does not work

Using vmap

Hi, I want to perform an vmap iteration over an array, but the "TracerArrayConversionError" appears. I want to ask if it is possible to use this kind of operation or not.
Thanks for your answer.

Recommend Projects