fix: ensure we properly indicate to jax that beta does not have grads for moffat and spergel

jax_galsim/moffat.py

@@ -72,7 +72,13 @@ def body(i, xcur):
     return jax.lax.fori_loop(0, 100, body, re, unroll=True)
 
 
-@implements(_galsim.Moffat)
+@implements(
+    _galsim.Moffat,
+    lax_description="""\
+The JAX implementation of the Moffat profile does not
+support autodiff with respect to the `beta` parameter.
+""",
+)
 @register_pytree_node_class
 class Moffat(GSObject):
     _is_axisymmetric = True
@@ -89,6 +95,9 @@ def __init__(
         flux=1.0,
         gsparams=None,
     ):
+        # beta doesn't have gradients to mark that right away
+        beta = jax.lax.stop_gradient(beta)
+
         # notice that trunc==0. means no truncated Moffat.
         # let define beta_thr a threshold to trigger the truncature
         self._beta_thr = 1.1

jax_galsim/spergel.py

@@ -255,6 +255,9 @@ def _spergel_hlr_pade(x):
             \hat{I}(k) = flux / (1 + (k r_0)^2)^{1+\nu}
 
     where :math:`r_0` is the ``scale_radius``, and :math: `\nu` mandatory to be in [-0.85,4.0]
+
+    The JAX Spergel profile does not support differentiation with respect to the :math:`\nu`
+    parameter.
     """,
 )
 @register_pytree_node_class
@@ -275,6 +278,8 @@ def __init__(
         flux=1.0,
         gsparams=None,
     ):
+        nu = jax.lax.stop_gradient(nu)
+
         # Parse the radius options
         if half_light_radius is not None:
             if scale_radius is not None: