Repository: sympy/sympy

Issue:
`JaxPrinter.printmethod` is an invalid identifier
The code printers allow custom printers to be defined for classes using specifically-named class methods. For example, the `NumPyPrinter` defines it's `printmethod` as `_numpycode`. However, it does this dynamically using the str contained within the `_module` attribute.

https://github.com/sympy/sympy/blob/1d9d3aa1aeddcce2f50fb922f359739ba16586e8/sympy/printing/numpy.py#L36-L53

The `JaxPrinter` class subclasses `NumPyPrinter` and so inherits its behaviour for setting the `printmethod` attribute. For the `JaxPrinter`, the `_module` is named `jax.numpy`. This results in `printmethod` being set to `_jax.numpycode`, which is an invalid identifier. This makes it impossible to define custom JAX printing for custom classes.

The solution is to reallocate the `printmethod` for JAX. I propose just making this `_jaxcode`.

Discussion:

Patch:
diff --git a/doc/src/explanation/active-deprecations.md b/doc/src/explanation/active-deprecations.md
index e166f68b9c..d1cc08ef69 100644
--- a/doc/src/explanation/active-deprecations.md
+++ b/doc/src/explanation/active-deprecations.md
@@ -637,14 +637,14 @@ Plotly, Mayavi, K3D only require lists of coordinates), this has been moved
 inside the `MatplotlibBackend` class.
 
 Note that previously, the method
-{meth}`~sympy.plotting.plot.LineOver1DRangeSeries.get_points` always returned
+{meth}`~sympy.plotting.plot.Parametric2DLineSeries.get_points` always returned
 uniformly sampled points, which meant that some functions were not plotted
 correctly when using `get_points()` to plot with Matplotlib.
 
 To avoid this problem, the method `get_segments()` could be used, which used
 adaptive sampling and which could be used with Matplotlib's `LineCollection`.
 However, this has been changed, and now `get_points()` can also use adaptive
-sampling. The {meth}`~sympy.plotting.plot.Line2DBaseSeries.get_data()` method can also be used.
+sampling. The {meth}`~.get_data()` method can also be used.
 
 
 (deprecated-physics-mdft)=
diff --git a/sympy/ntheory/factor_.py b/sympy/ntheory/factor_.py
index d5d0d02019..40e0b38b08 100644
--- a/sympy/ntheory/factor_.py
+++ b/sympy/ntheory/factor_.py
@@ -8,6 +8,7 @@
 
 from sympy.core import sympify
 from sympy.core.containers import Dict
+from sympy.core.evalf import bitcount
 from sympy.core.expr import Expr
 from sympy.core.function import Function
 from sympy.core.logic import fuzzy_and
@@ -208,12 +209,6 @@ def trailing(n):
     7
     >>> trailing(63)
     0
-
-    See Also
-    ========
-
-    multiplicity
-
     """
     n = abs(int(n))
     if not n:
@@ -222,28 +217,32 @@ def trailing(n):
     if low_byte:
         return small_trailing[low_byte]
 
-    t = 8
-    n >>= 8
     # 2**m is quick for z up through 2**30
-    z = n.bit_length() - 1
-    if n == 1 << z:
-        return z + t
+    z = bitcount(n) - 1
+    if isinstance(z, SYMPY_INTS):
+        if n == 1 << z:
+            return z
 
     if z < 300:
         # fixed 8-byte reduction
+        t = 8
+        n >>= 8
         while not n & 0xff:
             n >>= 8
             t += 8
-    else:
-        # binary reduction important when there might be a large
-        # number of trailing 0s
-        p = z >> 1
-        while not n & 0xff:
-            while n & ((1 << p) - 1):
-                p >>= 1
+        return t + small_trailing[n & 0xff]
+
+    # binary reduction important when there might be a large
+    # number of trailing 0s
+    t = 0
+    p = 8
+    while not n & 1:
+        while not n & ((1 << p) - 1):
             n >>= p
             t += p
-    return t + small_trailing[n & 0xff]
+            p *= 2
+        p //= 2
+    return t
 
 
 def multiplicity(p, n):
@@ -273,11 +272,6 @@ def multiplicity(p, n):
     >>> _ == multiplicity_in_factorial(p, n)
     True
 
-    See Also
-    ========
-
-    trailing
-
     """
     try:
         p, n = as_int(p), as_int(n)
diff --git a/sympy/ntheory/residue_ntheory.py b/sympy/ntheory/residue_ntheory.py
index 05dc426ccd..7341289e98 100644
--- a/sympy/ntheory/residue_ntheory.py
+++ b/sympy/ntheory/residue_ntheory.py
@@ -1616,7 +1616,7 @@ def polynomial_congruence(expr, m):
     Parameters
     ==========
 
-    expr : integer coefficient polynomial
+    coefficients : Coefficients of the Polynomial
     m : positive integer
 
     Examples
diff --git a/sympy/plotting/plot.py b/sympy/plotting/plot.py
index 5cd5987569..d783b59dac 100644
--- a/sympy/plotting/plot.py
+++ b/sympy/plotting/plot.py
@@ -809,7 +809,7 @@ def get_data(self):
                 List of z-coordinates in case of Parametric3DLineSeries
         """
         np = import_module('numpy')
-        points = self._get_data_helper()
+        points = self.get_points()
         if self.steps is True:
             if len(points) == 2:
                 x = np.array((points[0], points[0])).T.flatten()[1:]
@@ -897,13 +897,10 @@ def __str__(self):
             str(self.expr), str(self.var), str((self.start, self.end)))
 
     def get_points(self):
-        """Return lists of coordinates for plotting. Depending on the
+        """ Return lists of coordinates for plotting. Depending on the
         ``adaptive`` option, this function will either use an adaptive algorithm
         or it will uniformly sample the expression over the provided range.
 
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-
         Returns
         =======
             x : list
@@ -911,16 +908,12 @@ def get_points(self):
 
             y : list
                 List of y-coordinates
-        """
-        return self._get_data_helper()
 
-    def _get_data_helper(self):
-        if self.only_integers or not self.adaptive:
-            return self._uniform_sampling()
-        return self._adaptive_sampling()
 
-    def _adaptive_sampling(self):
-        """The adaptive sampling is done by recursively checking if three
+        Explanation
+        ===========
+
+        The adaptive sampling is done by recursively checking if three
         points are almost collinear. If they are not collinear, then more
         points are added between those points.
 
@@ -929,69 +922,73 @@ def _adaptive_sampling(self):
 
         .. [1] Adaptive polygonal approximation of parametric curves,
                Luiz Henrique de Figueiredo.
-        """
-        f = lambdify([self.var], self.expr)
-        x_coords = []
-        y_coords = []
-        np = import_module('numpy')
-        def sample(p, q, depth):
-            """ Samples recursively if three points are almost collinear.
-            For depth < 6, points are added irrespective of whether they
-            satisfy the collinearity condition or not. The maximum depth
-            allowed is 12.
-            """
-            # Randomly sample to avoid aliasing.
-            random = 0.45 + np.random.rand() * 0.1
-            if self.xscale == 'log':
-                xnew = 10**(np.log10(p[0]) + random * (np.log10(q[0]) -
-                                                        np.log10(p[0])))
-            else:
-                xnew = p[0] + random * (q[0] - p[0])
-            ynew = f(xnew)
-            new_point = np.array([xnew, ynew])
-
-            # Maximum depth
-            if depth > self.depth:
-                x_coords.append(q[0])
-                y_coords.append(q[1])
-
-            # Sample to depth of 6 (whether the line is flat or not)
-            # without using linspace (to avoid aliasing).
-            elif depth < 6:
-                sample(p, new_point, depth + 1)
-                sample(new_point, q, depth + 1)
 
-            # Sample ten points if complex values are encountered
-            # at both ends. If there is a real value in between, then
-            # sample those points further.
-            elif p[1] is None and q[1] is None:
+        """
+        if self.only_integers or not self.adaptive:
+            return self._uniform_sampling()
+        else:
+            f = lambdify([self.var], self.expr)
+            x_coords = []
+            y_coords = []
+            np = import_module('numpy')
+            def sample(p, q, depth):
+                """ Samples recursively if three points are almost collinear.
+                For depth < 6, points are added irrespective of whether they
+                satisfy the collinearity condition or not. The maximum depth
+                allowed is 12.
+                """
+                # Randomly sample to avoid aliasing.
+                random = 0.45 + np.random.rand() * 0.1
                 if self.xscale == 'log':
-                    xarray = np.logspace(p[0], q[0], 10)
+                    xnew = 10**(np.log10(p[0]) + random * (np.log10(q[0]) -
+                                                           np.log10(p[0])))
                 else:
-                    xarray = np.linspace(p[0], q[0], 10)
-                yarray = list(map(f, xarray))
-                if not all(y is None for y in yarray):
-                    for i in range(len(yarray) - 1):
-                        if not (yarray[i] is None and yarray[i + 1] is None):
-                            sample([xarray[i], yarray[i]],
-                                [xarray[i + 1], yarray[i + 1]], depth + 1)
-
-            # Sample further if one of the end points in None (i.e. a
-            # complex value) or the three points are not almost collinear.
-            elif (p[1] is None or q[1] is None or new_point[1] is None
-                    or not flat(p, new_point, q)):
-                sample(p, new_point, depth + 1)
-                sample(new_point, q, depth + 1)
-            else:
-                x_coords.append(q[0])
-                y_coords.append(q[1])
+                    xnew = p[0] + random * (q[0] - p[0])
+                ynew = f(xnew)
+                new_point = np.array([xnew, ynew])
+
+                # Maximum depth
+                if depth > self.depth:
+                    x_coords.append(q[0])
+                    y_coords.append(q[1])
+
+                # Sample irrespective of whether the line is flat till the
+                # depth of 6. We are not using linspace to avoid aliasing.
+                elif depth < 6:
+                    sample(p, new_point, depth + 1)
+                    sample(new_point, q, depth + 1)
+
+                # Sample ten points if complex values are encountered
+                # at both ends. If there is a real value in between, then
+                # sample those points further.
+                elif p[1] is None and q[1] is None:
+                    if self.xscale == 'log':
+                        xarray = np.logspace(p[0], q[0], 10)
+                    else:
+                        xarray = np.linspace(p[0], q[0], 10)
+                    yarray = list(map(f, xarray))
+                    if not all(y is None for y in yarray):
+                        for i in range(len(yarray) - 1):
+                            if not (yarray[i] is None and yarray[i + 1] is None):
+                                sample([xarray[i], yarray[i]],
+                                    [xarray[i + 1], yarray[i + 1]], depth + 1)
+
+                # Sample further if one of the end points in None (i.e. a
+                # complex value) or the three points are not almost collinear.
+                elif (p[1] is None or q[1] is None or new_point[1] is None
+                        or not flat(p, new_point, q)):
+                    sample(p, new_point, depth + 1)
+                    sample(new_point, q, depth + 1)
+                else:
+                    x_coords.append(q[0])
+                    y_coords.append(q[1])
 
-        f_start = f(self.start)
-        f_end = f(self.end)
-        x_coords.append(self.start)
-        y_coords.append(f_start)
-        sample(np.array([self.start, f_start]),
-                np.array([self.end, f_end]), 0)
+            f_start = f(self.start)
+            f_end = f(self.end)
+            x_coords.append(self.start)
+            y_coords.append(f_start)
+            sample(np.array([self.start, f_start]),
+                   np.array([self.end, f_end]), 0)
 
         return (x_coords, y_coords)
 
@@ -1043,14 +1040,19 @@ def get_parameter_points(self):
         np = import_module('numpy')
         return np.linspace(self.start, self.end, num=self.nb_of_points)
 
+    def _uniform_sampling(self):
+        param = self.get_parameter_points()
+        fx = vectorized_lambdify([self.var], self.expr_x)
+        fy = vectorized_lambdify([self.var], self.expr_y)
+        list_x = fx(param)
+        list_y = fy(param)
+        return (list_x, list_y)
+
     def get_points(self):
         """ Return lists of coordinates for plotting. Depending on the
         ``adaptive`` option, this function will either use an adaptive algorithm
         or it will uniformly sample the expression over the provided range.
 
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-
         Returns
         =======
             x : list
@@ -1059,24 +1061,11 @@ def get_points(self):
             y : list
                 List of y-coordinates
 
-        """
-        return self._get_data_helper()[:-1]
-
-    def _get_data_helper(self):
-        if not self.adaptive:
-            return self._uniform_sampling()
-        return self._adaptive_sampling()
 
-    def _uniform_sampling(self):
-        param = self.get_parameter_points()
-        fx = vectorized_lambdify([self.var], self.expr_x)
-        fy = vectorized_lambdify([self.var], self.expr_y)
-        list_x = fx(param)
-        list_y = fy(param)
-        return (list_x, list_y, param)
+        Explanation
+        ===========
 
-    def _adaptive_sampling(self):
-        """The adaptive sampling is done by recursively checking if three
+        The adaptive sampling is done by recursively checking if three
         points are almost collinear. If they are not collinear, then more
         points are added between those points.
 
@@ -1085,12 +1074,15 @@ def _adaptive_sampling(self):
 
         .. [1] Adaptive polygonal approximation of parametric curves,
             Luiz Henrique de Figueiredo.
+
         """
+        if not self.adaptive:
+            return self._uniform_sampling()
+
         f_x = lambdify([self.var], self.expr_x)
         f_y = lambdify([self.var], self.expr_y)
         x_coords = []
         y_coords = []
-        param = []
 
         def sample(param_p, param_q, p, q, depth):
             """ Samples recursively if three points are almost collinear.
@@ -1110,7 +1102,6 @@ def sample(param_p, param_q, p, q, depth):
             if depth > self.depth:
                 x_coords.append(q[0])
                 y_coords.append(q[1])
-                param.append(param_p)
 
             # Sample irrespective of whether the line is flat till the
             # depth of 6. We are not using linspace to avoid aliasing.
@@ -1146,7 +1137,6 @@ def sample(param_p, param_q, p, q, depth):
             else:
                 x_coords.append(q[0])
                 y_coords.append(q[1])
-                param.append(param_p)
 
         f_start_x = f_x(self.start)
         f_start_y = f_y(self.start)
@@ -1156,10 +1146,9 @@ def sample(param_p, param_q, p, q, depth):
         end = [f_end_x, f_end_y]
         x_coords.append(f_start_x)
         y_coords.append(f_start_y)
-        param.append(start)
         sample(self.start, self.end, start, end, 0)
 
-        return x_coords, y_coords, param
+        return x_coords, y_coords
 
 
 ### 3D lines
@@ -1208,23 +1197,6 @@ def get_parameter_points(self):
         return np.linspace(self.start, self.end, num=self.nb_of_points)
 
     def get_points(self):
-        """Return the x,y,z coordinates for plotting the line.
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-        """
-        return self.get_data()[:-1]
-
-    def get_data(self):
-        """Return coordinates for plotting the line.
-
-        Returns
-        =======
-
-        x: np.ndarray
-        y: np.ndarray
-        z: np.ndarray
-        param : np.ndarray
-        """
         np = import_module('numpy')
         param = self.get_parameter_points()
         fx = vectorized_lambdify([self.var], self.expr_x)
@@ -1246,7 +1218,7 @@ def get_data(self):
         self._xlim = (np.amin(list_x), np.amax(list_x))
         self._ylim = (np.amin(list_y), np.amax(list_y))
         self._zlim = (np.amin(list_z), np.amax(list_z))
-        return list_x, list_y, list_z, param
+        return list_x, list_y, list_z
 
 
 ### Surfaces
@@ -1314,24 +1286,6 @@ def __str__(self):
                     str((self.start_y, self.end_y)))
 
     def get_meshes(self):
-        """Return the x,y,z coordinates for plotting the surface.
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-        """
-        return self.get_data()
-
-    def get_data(self):
-        """Return arrays of coordinates for plotting.
-
-        Returns
-        =======
-        mesh_x : np.ndarray
-            Discretized x-domain.
-        mesh_y : np.ndarray
-            Discretized y-domain.
-        mesh_z : np.ndarray
-            Results of the evaluation.
-        """
         np = import_module('numpy')
         mesh_x, mesh_y = np.meshgrid(np.linspace(self.start_x, self.end_x,
                                                  num=self.nb_of_points_x),
@@ -1387,28 +1341,6 @@ def get_parameter_meshes(self):
                                        num=self.nb_of_points_v))
 
     def get_meshes(self):
-        """Return the x,y,z coordinates for plotting the surface.
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-        """
-        return self.get_data()[:3]
-
-    def get_data(self):
-        """Return arrays of coordinates for plotting.
-
-        Returns
-        =======
-        mesh_x : np.ndarray
-            Discretized x-domain.
-        mesh_y : np.ndarray
-            Discretized y-domain.
-        mesh_z : np.ndarray
-            Results of the evaluation.
-        mesh_u : np.ndarray
-            Discretized u range.
-        mesh_v : np.ndarray
-            Discretized v range.
-        """
         np = import_module('numpy')
 
         mesh_u, mesh_v = self.get_parameter_meshes()
@@ -1432,7 +1364,7 @@ def get_data(self):
         self._ylim = (np.amin(mesh_y), np.amax(mesh_y))
         self._zlim = (np.amin(mesh_z), np.amax(mesh_z))
 
-        return mesh_x, mesh_y, mesh_z, mesh_u, mesh_v
+        return mesh_x, mesh_y, mesh_z
 
 
 ### Contours
@@ -1470,24 +1402,6 @@ def __str__(self):
                     str((self.start_y, self.end_y)))
 
     def get_meshes(self):
-        """Return the x,y,z coordinates for plotting the surface.
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-        """
-        return self.get_data()[:3]
-
-    def get_data(self):
-        """Return arrays of coordinates for plotting.
-
-        Returns
-        =======
-        mesh_x : np.ndarray
-            Discretized x-domain.
-        mesh_y : np.ndarray
-            Discretized y-domain.
-        mesh_z : np.ndarray
-            Results of the evaluation.
-        """
         np = import_module('numpy')
         mesh_x, mesh_y = np.meshgrid(np.linspace(self.start_x, self.end_x,
                                                  num=self.nb_of_points_x),
@@ -1642,10 +1556,7 @@ def _process_series(self, series, ax):
         for s in series:
             # Create the collections
             if s.is_2Dline:
-                if s.is_parametric:
-                    x, y, param = s.get_data()
-                else:
-                    x, y = s.get_data()
+                x, y = s.get_data()
                 if (isinstance(s.line_color, (int, float)) or
                         callable(s.line_color)):
                     segments = self.get_segments(x, y)
@@ -1656,9 +1567,9 @@ def _process_series(self, series, ax):
                     lbl = _str_or_latex(s.label)
                     line, = ax.plot(x, y, label=lbl, color=s.line_color)
             elif s.is_contour:
-                ax.contour(*s.get_data())
+                ax.contour(*s.get_meshes())
             elif s.is_3Dline:
-                x, y, z, param = s.get_data()
+                x, y, z = s.get_data()
                 if (isinstance(s.line_color, (int, float)) or
                         callable(s.line_color)):
                     art3d = mpl_toolkits.mplot3d.art3d
@@ -1674,10 +1585,7 @@ def _process_series(self, series, ax):
                 ylims.append(s._ylim)
                 zlims.append(s._zlim)
             elif s.is_3Dsurface:
-                if s.is_parametric:
-                    x, y, z, u, v = s.get_data()
-                else:
-                    x, y, z = s.get_data()
+                x, y, z = s.get_meshes()
                 collection = ax.plot_surface(x, y, z,
                     cmap=getattr(self.cm, 'viridis', self.cm.jet),
                     rstride=1, cstride=1, linewidth=0.1)
@@ -1692,7 +1600,7 @@ def _process_series(self, series, ax):
                 ylims.append(s._ylim)
                 zlims.append(s._zlim)
             elif s.is_implicit:
-                points = s.get_data()
+                points = s.get_raster()
                 if len(points) == 2:
                     # interval math plotting
                     x, y = _matplotlib_list(points[0])
diff --git a/sympy/plotting/plot_implicit.py b/sympy/plotting/plot_implicit.py
index 019e99d408..68a6ecf92a 100644
--- a/sympy/plotting/plot_implicit.py
+++ b/sympy/plotting/plot_implicit.py
@@ -79,30 +79,7 @@ def __str__(self):
                     str(self.var_y),
                     str((self.start_y, self.end_y)))
 
-    def get_data(self):
-        """Returns numerical data.
-
-        Returns
-        =======
-
-        If the series is evaluated with the `adaptive=True` it returns:
-
-        interval_list : list
-            List of bounding rectangular intervals to be postprocessed and
-            eventually used with Matplotlib's ``fill`` command.
-        dummy : str
-            A string containing ``"fill"``.
-
-        Otherwise, it returns 2D numpy arrays to be used with Matplotlib's
-        ``contour`` or ``contourf`` commands:
-
-        x_array : np.ndarray
-        y_array : np.ndarray
-        z_array : np.ndarray
-        plot_type : str
-            A string specifying which plot command to use, ``"contour"``
-            or ``"contourf"``.
-        """
+    def get_raster(self):
         func = experimental_lambdify((self.var_x, self.var_y), self.expr,
                                     use_interval=True)
         xinterval = interval(self.start_x, self.end_x)
@@ -123,14 +100,6 @@ def get_data(self):
         else:
             return self._get_meshes_grid()
 
-    def get_raster(self):
-        """Returns numerical data.
-
-        This function is available for back-compatibility purposes. Consider
-        using ``get_data()`` instead.
-        """
-        return self.get_data()
-
     def _get_raster_interval(self, func):
         """ Uses interval math to adaptively mesh and obtain the plot"""
         k = self.depth
diff --git a/sympy/printing/numpy.py b/sympy/printing/numpy.py
index d686385d0a..761363ee5f 100644
--- a/sympy/printing/numpy.py
+++ b/sympy/printing/numpy.py
@@ -483,6 +483,7 @@ class JaxPrinter(NumPyPrinter):
 
     def __init__(self, settings=None):
         super().__init__(settings=settings)
+        self.printmethod = '_jaxcode'
 
     # These need specific override to allow for the lack of "jax.numpy.reduce"
     def _print_And(self, expr):


Test patch:
diff --git a/sympy/ntheory/primetest.py b/sympy/ntheory/primetest.py
index c3216ea562..58ce574d8c 100644
--- a/sympy/ntheory/primetest.py
+++ b/sympy/ntheory/primetest.py
@@ -3,10 +3,8 @@
 
 """
 
-from itertools import count
-
 from sympy.core.sympify import sympify
-from sympy.external.gmpy import HAS_GMPY, jacobi, is_square as gmpy_is_square
+from sympy.external.gmpy import HAS_GMPY, gcd, jacobi, is_square as gmpy_is_square
 from sympy.utilities.misc import as_int
 
 from mpmath.libmp import bitcount as _bitlength
@@ -241,100 +239,47 @@ def _lucas_sequence(n, P, Q, k):
 
 
 def _lucas_selfridge_params(n):
-    """Calculates the Selfridge parameters (D, P, Q) for n.
-
-    Explanation
-    ===========
-
-    The Lucas compositeness test checks whether n is a prime number.
-    The test can be run with arbitrary parameters ``P`` and ``Q``, which also change the performance of the test.
-    So, which parameters are most effective for running the Lucas compositeness test?
-    As an algorithm for determining ``P`` and ``Q``, Selfridge proposed method A [1]_ page 1401
-    (Since two methods were proposed, referred to simply as A and B in the paper,
-    we will refer to one of them as "method A").
-
-    method A fixes ``P = 1``. Then, ``D`` defined by ``D = P**2 - 4Q`` is varied from 5, -7, 9, -11, 13, and so on,
-    with the first ``D`` being ``jacobi(D, n) == -1``. Once ``D`` is determined,
-    ``Q`` is determined to be ``(P**2 - D)//4``. Here is an implementation of the method A.
-
-    Parameters
-    ==========
-
-    n : int
-        positive odd integer
-
-    Returns
-    =======
-
-    D, P, Q: Selfridge parameters for Lucas compositeness test.
-             ``(0, 0, 0)`` if we find a nontrivial divisor of ``n``.
-
-    Examples
-    ========
-
-    >>> from sympy.ntheory.primetest import _lucas_selfridge_params
-    >>> _lucas_selfridge_params(101)
-    (-7, 1, 2)
-    >>> _lucas_selfridge_params(15)
-    (0, 0, 0)
+    """Calculates the Selfridge parameters (D, P, Q) for n.  This is
+       method A from page 1401 of Baillie and Wagstaff.
 
     References
     ==========
-
-    .. [1] Robert Baillie, Samuel S. Wagstaff, Lucas Pseudoprimes,
-           Math. Comp. Vol 35, Number 152 (1980), pp. 1391-1417,
-           https://doi.org/10.1090%2FS0025-5718-1980-0583518-6
+    .. [1] "Lucas Pseudoprimes", Baillie and Wagstaff, 1980.
            http://mpqs.free.fr/LucasPseudoprimes.pdf
-
     """
-    for D in count(5, 2):
-        if D & 2: # if D % 4 == 3
-            D = -D
-        j = jacobi(D, n)
-        if j == -1:
-            return (D, 1, (1 - D)//4)
-        elif j == 0 and D % n:
+    D = 5
+    while True:
+        g = gcd(D, n)
+        if g > 1 and g != n:
             return (0, 0, 0)
+        if jacobi(D, n) == -1:
+            break
+        if D > 0:
+          D = -D - 2
+        else:
+          D = -D + 2
+    return _int_tuple(D, 1, (1 - D)/4)
 
 
 def _lucas_extrastrong_params(n):
     """Calculates the "extra strong" parameters (D, P, Q) for n.
 
-    Parameters
-    ==========
-
-    n : int
-        positive odd integer
-
-    Returns
-    =======
-
-    D, P, Q: "extra strong" parameters.
-             ``(0, 0, 0)`` if we find a nontrivial divisor of ``n``.
-
-    Examples
-    ========
-
-    >>> from sympy.ntheory.primetest import _lucas_extrastrong_params
-    >>> _lucas_extrastrong_params(101)
-    (12, 4, 1)
-    >>> _lucas_extrastrong_params(15)
-    (0, 0, 0)
-
     References
     ==========
     .. [1] OEIS A217719: Extra Strong Lucas Pseudoprimes
            https://oeis.org/A217719
-    .. [2] https://en.wikipedia.org/wiki/Lucas_pseudoprime
-
+    .. [1] https://en.wikipedia.org/wiki/Lucas_pseudoprime
     """
-    for P in count(3):
-        D = P**2 - 4
-        j = jacobi(D, n)
-        if j == -1:
-            return (D, P, 1)
-        elif j == 0 and D % n:
+    P, Q, D = 3, 1, 5
+    while True:
+        g = gcd(D, n)
+        if g > 1 and g != n:
             return (0, 0, 0)
+        if jacobi(D, n) == -1:
+            break
+        P += 1
+        D = P*P - 4
+    return _int_tuple(D, P, Q)
 
 
 def is_lucas_prp(n):
@@ -346,13 +291,11 @@ def is_lucas_prp(n):
 
     References
     ==========
-    .. [1] Robert Baillie, Samuel S. Wagstaff, Lucas Pseudoprimes,
-           Math. Comp. Vol 35, Number 152 (1980), pp. 1391-1417,
-           https://doi.org/10.1090%2FS0025-5718-1980-0583518-6
-           http://mpqs.free.fr/LucasPseudoprimes.pdf
-    .. [2] OEIS A217120: Lucas Pseudoprimes
-           https://oeis.org/A217120
-    .. [3] https://en.wikipedia.org/wiki/Lucas_pseudoprime
+    - "Lucas Pseudoprimes", Baillie and Wagstaff, 1980.
+      http://mpqs.free.fr/LucasPseudoprimes.pdf
+    - OEIS A217120: Lucas Pseudoprimes
+      https://oeis.org/A217120
+    - https://en.wikipedia.org/wiki/Lucas_pseudoprime
 
     Examples
     ========
@@ -360,7 +303,7 @@ def is_lucas_prp(n):
     >>> from sympy.ntheory.primetest import isprime, is_lucas_prp
     >>> for i in range(10000):
     ...     if is_lucas_prp(i) and not isprime(i):
-    ...         print(i)
+    ...        print(i)
     323
     377
     1159
@@ -382,7 +325,8 @@ def is_lucas_prp(n):
     D, P, Q = _lucas_selfridge_params(n)
     if D == 0:
         return False
-    return _lucas_sequence(n, P, Q, n + 1)[0] == 0
+    U, V, Qk = _lucas_sequence(n, P, Q, n+1)
+    return U == 0
 
 
 def is_strong_lucas_prp(n):
@@ -395,14 +339,12 @@ def is_strong_lucas_prp(n):
 
     References
     ==========
-    .. [1] Robert Baillie, Samuel S. Wagstaff, Lucas Pseudoprimes,
-           Math. Comp. Vol 35, Number 152 (1980), pp. 1391-1417,
-           https://doi.org/10.1090%2FS0025-5718-1980-0583518-6
-           http://mpqs.free.fr/LucasPseudoprimes.pdf
-    .. [2] OEIS A217255: Strong Lucas Pseudoprimes
-           https://oeis.org/A217255
-    .. [3] https://en.wikipedia.org/wiki/Lucas_pseudoprime
-    .. [4] https://en.wikipedia.org/wiki/Baillie-PSW_primality_test
+    - "Lucas Pseudoprimes", Baillie and Wagstaff, 1980.
+      http://mpqs.free.fr/LucasPseudoprimes.pdf
+    - OEIS A217255: Strong Lucas Pseudoprimes
+      https://oeis.org/A217255
+    - https://en.wikipedia.org/wiki/Lucas_pseudoprime
+    - https://en.wikipedia.org/wiki/Baillie-PSW_primality_test
 
     Examples
     ========
@@ -432,13 +374,13 @@ def is_strong_lucas_prp(n):
 
     # remove powers of 2 from n+1 (= k * 2**s)
     s = trailing(n + 1)
-    k = (n + 1) >> s
+    k = (n+1) >> s
 
     U, V, Qk = _lucas_sequence(n, P, Q, k)
 
     if U == 0 or V == 0:
         return True
-    for _ in range(1, s):
+    for r in range(1, s):
         V = (V*V - 2*Qk) % n
         if V == 0:
             return True
@@ -452,7 +394,7 @@ def is_extra_strong_lucas_prp(n):
     prime.
 
     The parameters are selected using P = 3, Q = 1, then incrementing P until
-    (D|n) == -1.  The test itself is as defined in [1]_, from the
+    (D|n) == -1.  The test itself is as defined in Grantham 2000, from the
     Mo and Jones preprint.  The parameter selection and test are the same as
     used in OEIS A217719, Perl's Math::Prime::Util, and the Lucas pseudoprime
     page on Wikipedia.
@@ -465,12 +407,11 @@ def is_extra_strong_lucas_prp(n):
 
     References
     ==========
-    .. [1] Jon Grantham, Frobenius Pseudoprimes,
-           Math. Comp. Vol 70, Number 234 (2001), pp. 873-891,
-           https://doi.org/10.1090%2FS0025-5718-00-01197-2
-    .. [2] OEIS A217719: Extra Strong Lucas Pseudoprimes
-           https://oeis.org/A217719
-    .. [3] https://en.wikipedia.org/wiki/Lucas_pseudoprime
+    - "Frobenius Pseudoprimes", Jon Grantham, 2000.
+      https://www.ams.org/journals/mcom/2001-70-234/S0025-5718-00-01197-2/
+    - OEIS A217719: Extra Strong Lucas Pseudoprimes
+      https://oeis.org/A217719
+    - https://en.wikipedia.org/wiki/Lucas_pseudoprime
 
     Examples
     ========
@@ -506,13 +447,13 @@ def is_extra_strong_lucas_prp(n):
 
     # remove powers of 2 from n+1 (= k * 2**s)
     s = trailing(n + 1)
-    k = (n + 1) >> s
+    k = (n+1) >> s
 
-    U, V, _ = _lucas_sequence(n, P, Q, k)
+    U, V, Qk = _lucas_sequence(n, P, Q, k)
 
     if U == 0 and (V == 2 or V == n - 2):
         return True
-    for _ in range(1, s):
+    for r in range(1, s):
         if V == 0:
             return True
         V = (V*V - 2) % n
diff --git a/sympy/ntheory/tests/test_factor_.py b/sympy/ntheory/tests/test_factor_.py
index 995062aae8..626dde0dec 100644
--- a/sympy/ntheory/tests/test_factor_.py
+++ b/sympy/ntheory/tests/test_factor_.py
@@ -76,9 +76,6 @@ def test_trailing_bitcount():
     for i in range(100):
         assert trailing(1 << i) == i
         assert trailing((1 << i) * 31337) == i
-    for i in range(500):
-        n = (1 << 500) + (1 << i)
-        assert trailing(n) == i
     assert trailing(1 << 1000001) == 1000001
     assert trailing((1 << 273956)*7**37) == 273956
     # issue 12709
diff --git a/sympy/ntheory/tests/test_primetest.py b/sympy/ntheory/tests/test_primetest.py
index cf91f05303..8817b645a0 100644
--- a/sympy/ntheory/tests/test_primetest.py
+++ b/sympy/ntheory/tests/test_primetest.py
@@ -1,5 +1,5 @@
 from sympy.ntheory.generate import Sieve, sieve
-from sympy.ntheory.primetest import (mr, _lucas_selfridge_params, _lucas_extrastrong_params, is_lucas_prp, is_square,
+from sympy.ntheory.primetest import (mr, is_lucas_prp, is_square,
                                      is_strong_lucas_prp, is_extra_strong_lucas_prp, isprime, is_euler_pseudoprime,
                                      is_gaussian_prime)
 
@@ -25,27 +25,6 @@ def test_euler_pseudoprimes():
     assert is_euler_pseudoprime(8421, 29) == False
     assert is_euler_pseudoprime(9997, 19) == True
 
-
-def test_lucas_selfridge_params():
-    assert _lucas_selfridge_params(3) == (5, 1, -1)
-    assert _lucas_selfridge_params(5) == (-7, 1, 2)
-    assert _lucas_selfridge_params(7) == (5, 1, -1)
-    assert _lucas_selfridge_params(9) == (0, 0, 0)
-    assert _lucas_selfridge_params(11) == (13, 1, -3)
-    assert _lucas_selfridge_params(19) == (-7, 1, 2)
-    assert _lucas_selfridge_params(29) == (-11, 1, 3)
-
-
-def test_lucas_extrastrong_params():
-    assert _lucas_extrastrong_params(3) == (5, 3, 1)
-    assert _lucas_extrastrong_params(5) == (12, 4, 1)
-    assert _lucas_extrastrong_params(7) == (5, 3, 1)
-    assert _lucas_extrastrong_params(9) == (0, 0, 0)
-    assert _lucas_extrastrong_params(11) == (21, 5, 1)
-    assert _lucas_extrastrong_params(59) == (32, 6, 1)
-    assert _lucas_extrastrong_params(479) == (117, 11, 1)
-
-
 def test_is_extra_strong_lucas_prp():
     assert is_extra_strong_lucas_prp(4) == False
     assert is_extra_strong_lucas_prp(989) == True
diff --git a/sympy/plotting/tests/test_plot.py b/sympy/plotting/tests/test_plot.py
index e642cbfe83..aa86d51824 100644
--- a/sympy/plotting/tests/test_plot.py
+++ b/sympy/plotting/tests/test_plot.py
@@ -802,31 +802,3 @@ def test_deprecated_markers_annotations_rectangles_fill():
     with warns_deprecated_sympy():
         p.rectangles = [{"xy": (0, 0), "width": 5, "height": 1}]
     assert len(p._series) == 5
-
-
-def test_back_compatibility():
-    if not matplotlib:
-        skip("Matplotlib not the default backend")
-
-    x = Symbol('x')
-    y = Symbol('y')
-    p = plot(sin(x), adaptive=False, n=5)
-    assert len(p[0].get_points()) == 2
-    assert len(p[0].get_data()) == 2
-    p = plot_parametric(cos(x), sin(x), (x, 0, 2), adaptive=False, n=5)
-    assert len(p[0].get_points()) == 2
-    assert len(p[0].get_data()) == 3
-    p = plot3d_parametric_line(cos(x), sin(x), x, (x, 0, 2),
-        adaptive=False, n=5)
-    assert len(p[0].get_points()) == 3
-    assert len(p[0].get_data()) == 4
-    p = plot3d(cos(x**2 + y**2), (x, -pi, pi), (y, -pi, pi), n=5)
-    assert len(p[0].get_meshes()) == 3
-    assert len(p[0].get_data()) == 3
-    p = plot_contour(cos(x**2 + y**2), (x, -pi, pi), (y, -pi, pi), n=5)
-    assert len(p[0].get_meshes()) == 3
-    assert len(p[0].get_data()) == 3
-    p = plot3d_parametric_surface(x * cos(y), x * sin(y), x * cos(4 * y) / 2,
-        (x, 0, pi), (y, 0, 2*pi), n=5)
-    assert len(p[0].get_meshes()) == 3
-    assert len(p[0].get_data()) == 5
diff --git a/sympy/printing/tests/test_jax.py b/sympy/printing/tests/test_jax.py
index a8428d4faf..4a58b0bada 100644
--- a/sympy/printing/tests/test_jax.py
+++ b/sympy/printing/tests/test_jax.py
@@ -9,7 +9,7 @@
 from sympy.utilities.lambdify import lambdify
 
 from sympy.abc import x, i, j, a, b, c, d
-from sympy.core import Pow
+from sympy.core import Function, Pow, Symbol
 from sympy.codegen.matrix_nodes import MatrixSolve
 from sympy.codegen.numpy_nodes import logaddexp, logaddexp2
 from sympy.codegen.cfunctions import log1p, expm1, hypot, log10, exp2, log2, Sqrt
@@ -36,6 +36,7 @@
     deafult_float_info = jax.numpy.finfo(jax.numpy.array([]).dtype)
     JAX_DEFAULT_EPSILON = deafult_float_info.eps
 
+
 def test_jax_piecewise_regression():
     """
     NumPyPrinter needs to print Piecewise()'s choicelist as a list to avoid
@@ -330,10 +331,12 @@ def test_issue_17006():
     N = MatrixSymbol("M", n, n)
     raises(NotImplementedError, lambda: lambdify(N, N + Identity(n), 'jax'))
 
+
 def test_jax_array():
     assert JaxPrinter().doprint(Array(((1, 2), (3, 5)))) == 'jax.numpy.array([[1, 2], [3, 5]])'
     assert JaxPrinter().doprint(Array((1, 2))) == 'jax.numpy.array((1, 2))'
 
+
 def test_jax_known_funcs_consts():
     assert _jax_known_constants['NaN'] == 'jax.numpy.nan'
     assert _jax_known_constants['EulerGamma'] == 'jax.numpy.euler_gamma'
@@ -341,7 +344,27 @@ def test_jax_known_funcs_consts():
     assert _jax_known_functions['acos'] == 'jax.numpy.arccos'
     assert _jax_known_functions['log'] == 'jax.numpy.log'
 
+
 def test_jax_print_methods():
     prntr = JaxPrinter()
     assert hasattr(prntr, '_print_acos')
     assert hasattr(prntr, '_print_log')
+
+
+def test_jax_printmethod():
+    printer = JaxPrinter()
+    assert hasattr(printer, 'printmethod')
+    assert printer.printmethod == '_jaxcode'
+
+
+def test_jax_custom_print_method():
+
+    class expm1(Function):
+
+        def _jaxcode(self, printer):
+            x, = self.args
+            function = f'expm1({printer._print(x)})'
+            return printer._module_format(printer._module + '.' + function)
+
+    printer = JaxPrinter()
+    assert printer.doprint(expm1(Symbol('x'))) == 'jax.numpy.expm1(x)'