add doc for use_jax and test for use_jax

piff/basis_interp.py

@@ -417,6 +417,7 @@ class BasisPolynomial(BasisInterp):
                         and is somewhat slower (nearly a factor of 2); however, it is significantly
                         less susceptible to numerical errors from high condition matrices.
                         Therefore, it may be preferred for some use cases. [default: False]
+    :param use_jax:     Use JAX in _solve_direct compared to classic numpy/scipy. [default: False]
     :param logger:      A logger object for logging debug info. [default: None]
     """
     _type_name = 'BasisPolynomial'

tests/test_pixel.py

@@ -1025,7 +1025,8 @@ def test_single_image():
             },
             'interp' : {
                 'type' : 'BasisPolynomial',
-                'order' : order
+                'order' : order,
+                'use_jax' : True,
             },
         },
     }