From f3dd05a261dc2d616a1efab20081135eac5b415d Mon Sep 17 00:00:00 2001
From: Sam Schiavone <sam.schiavone@gmail.com>
Date: Fri, 17 Apr 2026 11:55:10 -0400
Subject: [PATCH 01/15] working on getting acb_theta_jet in Magma; need to
 rebuild using meson

---
 src/flint/types/acb_mat.pyi   |   4 ++
 src/flint/types/acb_mat.pyx   |  19 +++++++
 src/flint/types/acb_theta.pyi |   2 +
 src/flint/types/acb_theta.pyx | 101 ++++++++++++++++++++++++++++++++++
 4 files changed, 126 insertions(+)

diff --git a/src/flint/types/acb_mat.pyi b/src/flint/types/acb_mat.pyi
index 87a8736b..c497cae1 100644
--- a/src/flint/types/acb_mat.pyi
+++ b/src/flint/types/acb_mat.pyi
@@ -97,6 +97,10 @@ class acb_mat(flint_mat[acb]):
 
     def theta(self, z: acb_mat, square: bool = False) -> acb_mat: ...
 
+    def theta_jet(self, z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
+
+acb_theta_jet(zvec, nb_in, tau, ord, ab, all, square=square)
+
     def str(self, n: int = 0, radius: bool = True, more: bool = False, condense: int = 0) -> _str: ...
     def repr(self) -> _str: ...
     def __str__(self) -> _str: ...
diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index ae65fe71..7969fd34 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -839,3 +839,22 @@ cdef class acb_mat(flint_mat):
         except ImportError:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
         return acb_theta(z, tau, square=square)
+
+    def theta_jet(zvec, nb_in, tau, ord, square=False):
+        r"""
+        Computes Taylor approximation for the vector-valued Riemann theta function
+        `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares,
+        where `\tau` is given by ``self``.
+
+        This is a wrapper for :func:`.acb_theta_jet.acb_theta_jet`; see the
+        documentation for that method for details and examples.
+        This follows the same conventions of the C-function
+        `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
+        for the ordering of the theta characteristics.
+
+        """
+        try:
+            from .acb_theta import acb_theta_jet
+        except ImportError:
+            raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
+        return acb_theta_jet(zvec, nb_in, tau, ord, square=square)
diff --git a/src/flint/types/acb_theta.pyi b/src/flint/types/acb_theta.pyi
index 37f1e1e3..37b69b74 100644
--- a/src/flint/types/acb_theta.pyi
+++ b/src/flint/types/acb_theta.pyi
@@ -4,3 +4,5 @@ from flint.types.acb_mat import acb_mat
 
 
 def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
+
+def acb_theta_jet(z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index 80ed5df1..a50c211f 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -95,3 +95,104 @@ def acb_theta(acb_mat z, acb_mat tau, ulong square=False):
         res.append(r)
     _acb_vec_clear(theta, nb)
     return acb_mat([res])
+
+#def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
+    #    r"""
+    #    Computes derivatives of the vector valued Riemann theta function
+    #    `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares.
+    #
+    #    This is a wrapper for the C-function
+    #    `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
+    #    and it follows the same conventions for the ordering of the theta characteristics.
+    #
+    #    This should be used via the method :meth:`.acb_mat.theta`, explicitly ``tau.theta(z)``.
+    #    """
+    #    g = tau.nrows()
+    #    assert tau.ncols() == g
+    #    assert z.nrows() == g
+    #    assert z.ncols() == 1
+    #
+    #    # convert input
+    #    cdef acb_ptr zvec
+    #    zvec = _acb_vec_init(g)
+    #    cdef long i
+    #    for i in range(g):
+    #        acb_set(zvec + i, acb_mat_entry(z.val, i, 0))
+    #    cdef slong nb_in = 1 
+    #    cdef ulong ab = 0 
+    #    cdef ulong all = True
+    #
+    #    # initialize the output
+    #    cdef slong nb = 1 << (2 * g)
+    #    # 1. Calculate the length of the jet for one characteristic
+    #    # This is the number of multi-indices (alpha) such that |alpha| < ord
+    #    cdef slong nj = acb_theta_jet_len(g, ord)
+    #    cdef acb_ptr theta = _acb_vec_init(nb)
+    #
+    #    acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
+    #    _acb_vec_clear(zvec, g)
+    #    # copy the output
+    #    res = []
+    #    cdef acb r
+    #    for i in range(nb):
+    #        r = acb.__new__(acb)
+    #        acb_set(r.val, theta + i)
+    #        res.append(r)
+    #    _acb_vec_clear(theta, nb)
+    #    return acb_mat([res])
+
+def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
+    r"""
+    Corrected wrapper for acb_theta_jet to handle multivariate Taylor expansions.
+    """
+    cdef slong g = tau.nrows()
+    if g == 0:
+        return []
+
+    # 1. Calculate the length of the jet for one characteristic
+    # This is the number of multi-indices (alpha) such that |alpha| < ord
+    cdef slong nj = acb_theta_jet_len(g, ord)
+    
+    # 2. Total number of characteristics
+    cdef slong nb = 1 << (2 * g)
+    
+    # 3. Total number of acb elements to allocate
+    # FLINT stores nj coefficients for each of the nb characteristics
+    cdef slong total_size = nb * nj
+
+    # Convert input z to acb_ptr
+    cdef acb_ptr zvec = _acb_vec_init(g)
+    cdef slong i, j
+    for i in range(g):
+        acb_set(zvec + i, acb_mat_entry(z.val, i, 0))
+
+    # Parameters for characteristics
+    cdef slong nb_in = 1  # Number of input z vectors
+    cdef ulong ab = 0     # Base characteristic
+    cdef ulong all = True # Compute all 2^2g characteristics
+
+    # Initialize the output buffer
+    cdef acb_ptr theta = _acb_vec_init(total_size)
+
+    # Call the FLINT C function
+    # Note: Computes all partial derivatives up to total order 'ord'
+    acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
+
+    # 4. Copy the output into a structured format
+    # We return a list of lists: res[char_idx] = [coeff_0, coeff_1, ...]
+    res = []
+    cdef acb r
+    for i in range(nb):
+        char_jet = []
+        for j in range(nj):
+            r = acb.__new__(acb)
+            # Offset: characteristic index * length of one jet + coefficient index
+            acb_set(r.val, theta + (i * nj + j))
+            char_jet.append(r)
+        res.append(char_jet)
+
+    # Cleanup
+    _acb_vec_clear(zvec, g)
+    _acb_vec_clear(theta, total_size)
+
+    return res

From 75b6aabcaa4f94b431de240859b7a00fcd14affb Mon Sep 17 00:00:00 2001
From: JHanselman <hanselmanj@hotmail.com>
Date: Fri, 24 Apr 2026 15:43:15 +0200
Subject: [PATCH 02/15] Some bug-fixing

---
 src/flint/types/acb_mat.pyi   |  3 +-
 src/flint/types/acb_mat.pyx   |  8 +--
 src/flint/types/acb_theta.pyi |  2 +-
 src/flint/types/acb_theta.pyx | 96 ++++++++++++++++++-----------------
 4 files changed, 55 insertions(+), 54 deletions(-)

diff --git a/src/flint/types/acb_mat.pyi b/src/flint/types/acb_mat.pyi
index c497cae1..755071a4 100644
--- a/src/flint/types/acb_mat.pyi
+++ b/src/flint/types/acb_mat.pyi
@@ -97,9 +97,8 @@ class acb_mat(flint_mat[acb]):
 
     def theta(self, z: acb_mat, square: bool = False) -> acb_mat: ...
 
-    def theta_jet(self, z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
+    def theta_jets(self, z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
 
-acb_theta_jet(zvec, nb_in, tau, ord, ab, all, square=square)
 
     def str(self, n: int = 0, radius: bool = True, more: bool = False, condense: int = 0) -> _str: ...
     def repr(self) -> _str: ...
diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index 7969fd34..8397a589 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -838,9 +838,9 @@ cdef class acb_mat(flint_mat):
             from .acb_theta import acb_theta
         except ImportError:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
-        return acb_theta(z, tau, square=square)
+        return acb_theta(z, 2*tau, square=square)
 
-    def theta_jet(zvec, nb_in, tau, ord, square=False):
+    def theta_jets(tau, z, ord, square=False):
         r"""
         Computes Taylor approximation for the vector-valued Riemann theta function
         `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares,
@@ -854,7 +854,7 @@ cdef class acb_mat(flint_mat):
 
         """
         try:
-            from .acb_theta import acb_theta_jet
+            from .acb_theta import acb_theta_jets
         except ImportError:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
-        return acb_theta_jet(zvec, nb_in, tau, ord, square=square)
+        return acb_theta_jets(z, tau, ord, square=square)
diff --git a/src/flint/types/acb_theta.pyi b/src/flint/types/acb_theta.pyi
index 37b69b74..5d06b495 100644
--- a/src/flint/types/acb_theta.pyi
+++ b/src/flint/types/acb_theta.pyi
@@ -5,4 +5,4 @@ from flint.types.acb_mat import acb_mat
 
 def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
 
-def acb_theta_jet(z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
+def acb_theta_jets(z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index a50c211f..f615af89 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -96,62 +96,18 @@ def acb_theta(acb_mat z, acb_mat tau, ulong square=False):
     _acb_vec_clear(theta, nb)
     return acb_mat([res])
 
-#def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
-    #    r"""
-    #    Computes derivatives of the vector valued Riemann theta function
-    #    `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares.
-    #
-    #    This is a wrapper for the C-function
-    #    `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
-    #    and it follows the same conventions for the ordering of the theta characteristics.
-    #
-    #    This should be used via the method :meth:`.acb_mat.theta`, explicitly ``tau.theta(z)``.
-    #    """
-    #    g = tau.nrows()
-    #    assert tau.ncols() == g
-    #    assert z.nrows() == g
-    #    assert z.ncols() == 1
-    #
-    #    # convert input
-    #    cdef acb_ptr zvec
-    #    zvec = _acb_vec_init(g)
-    #    cdef long i
-    #    for i in range(g):
-    #        acb_set(zvec + i, acb_mat_entry(z.val, i, 0))
-    #    cdef slong nb_in = 1 
-    #    cdef ulong ab = 0 
-    #    cdef ulong all = True
-    #
-    #    # initialize the output
-    #    cdef slong nb = 1 << (2 * g)
-    #    # 1. Calculate the length of the jet for one characteristic
-    #    # This is the number of multi-indices (alpha) such that |alpha| < ord
-    #    cdef slong nj = acb_theta_jet_len(g, ord)
-    #    cdef acb_ptr theta = _acb_vec_init(nb)
-    #
-    #    acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
-    #    _acb_vec_clear(zvec, g)
-    #    # copy the output
-    #    res = []
-    #    cdef acb r
-    #    for i in range(nb):
-    #        r = acb.__new__(acb)
-    #        acb_set(r.val, theta + i)
-    #        res.append(r)
-    #    _acb_vec_clear(theta, nb)
-    #    return acb_mat([res])
 
-def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
+def acb_theta_jets(acb_mat z, acb_mat tau, slong ord, ulong square=False):
     r"""
     Corrected wrapper for acb_theta_jet to handle multivariate Taylor expansions.
     """
-    cdef slong g = tau.nrows()
+    g = tau.nrows()
     if g == 0:
         return []
 
     # 1. Calculate the length of the jet for one characteristic
     # This is the number of multi-indices (alpha) such that |alpha| < ord
-    cdef slong nj = acb_theta_jet_len(g, ord)
+    cdef slong nj = acb_theta_jet_nb(g, ord)
     
     # 2. Total number of characteristics
     cdef slong nb = 1 << (2 * g)
@@ -196,3 +152,49 @@ def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
     _acb_vec_clear(theta, total_size)
 
     return res
+
+
+#def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
+    #    r"""
+    #    Computes derivatives of the vector valued Riemann theta function
+    #    `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares.
+    #
+    #    This is a wrapper for the C-function
+    #    `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
+    #    and it follows the same conventions for the ordering of the theta characteristics.
+    #
+    #    This should be used via the method :meth:`.acb_mat.theta`, explicitly ``tau.theta(z)``.
+    #    """
+    #    g = tau.nrows()
+    #    assert tau.ncols() == g
+    #    assert z.nrows() == g
+    #    assert z.ncols() == 1
+    #
+    #    # convert input
+    #    cdef acb_ptr zvec
+    #    zvec = _acb_vec_init(g)
+    #    cdef long i
+    #    for i in range(g):
+    #        acb_set(zvec + i, acb_mat_entry(z.val, i, 0))
+    #    cdef slong nb_in = 1 
+    #    cdef ulong ab = 0 
+    #    cdef ulong all = True
+    #
+    #    # initialize the output
+    #    cdef slong nb = 1 << (2 * g)
+    #    # 1. Calculate the length of the jet for one characteristic
+    #    # This is the number of multi-indices (alpha) such that |alpha| < ord
+    #    cdef slong nj = acb_theta_jet_len(g, ord)
+    #    cdef acb_ptr theta = _acb_vec_init(nb)
+    #
+    #    acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
+    #    _acb_vec_clear(zvec, g)
+    #    # copy the output
+    #    res = []
+    #    cdef acb r
+    #    for i in range(nb):
+    #        r = acb.__new__(acb)
+    #        acb_set(r.val, theta + i)
+    #        res.append(r)
+    #    _acb_vec_clear(theta, nb)
+    #    return acb_mat([res])

From cd11057443f4c0832052ad22d40a23e6e94e6c91 Mon Sep 17 00:00:00 2001
From: JHanselman <hanselmanj@hotmail.com>
Date: Wed, 6 May 2026 16:35:21 +0200
Subject: [PATCH 03/15] Removed nonsense

---
 src/flint/types/acb_mat.pyx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index 8397a589..caf01287 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -838,7 +838,7 @@ cdef class acb_mat(flint_mat):
             from .acb_theta import acb_theta
         except ImportError:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
-        return acb_theta(z, 2*tau, square=square)
+        return acb_theta(z, tau, square=square)
 
     def theta_jets(tau, z, ord, square=False):
         r"""

From 22b4e5a2e43a97ee10d040e1f3cdbe11e0d14944 Mon Sep 17 00:00:00 2001
From: JHanselman <hanselmanj@hotmail.com>
Date: Wed, 13 May 2026 17:08:03 +0200
Subject: [PATCH 04/15] Add tests and docstring

---
 src/flint/test/test_acb_theta.py | 19 +++++++
 src/flint/types/acb_mat.pyx      |  4 +-
 src/flint/types/acb_theta.pyx    | 85 ++++++++++++--------------------
 3 files changed, 52 insertions(+), 56 deletions(-)

diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index 4ea6237a..c1470d8d 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -45,3 +45,22 @@ def test_acb_theta_shape_assertions() -> None:
 
     assert raises(lambda: acb_theta(object(), tau), TypeError)  # type: ignore[arg-type]
     assert raises(lambda: acb_theta(z, object()), TypeError)  # type: ignore[arg-type]
+
+def test_acb_theta_jets_basic() -> None:
+  if not _has_acb_theta():
+      return
+
+  from flint.types.acb_theta import acb_theta_jets
+
+  z = acb(1 + 1j)
+  tau = acb(1.25 + 3j)
+  zmat = acb_mat([[z]])
+  taumat = acb_mat([[tau]])
+  ord = 2
+
+  direct = acb_theta_jets(zmat, taumat, ord)
+  via_method = taumat.theta_jets(zmat, ord)
+  assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12)
+  assert direct.nrows() == 4
+  assert direct.ncols() == 3
+
diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index caf01287..0b67f2a4 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -840,7 +840,7 @@ cdef class acb_mat(flint_mat):
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
         return acb_theta(z, tau, square=square)
 
-    def theta_jets(tau, z, ord, square=False):
+    def theta_jets(tau, z, ord):
         r"""
         Computes Taylor approximation for the vector-valued Riemann theta function
         `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares,
@@ -857,4 +857,4 @@ cdef class acb_mat(flint_mat):
             from .acb_theta import acb_theta_jets
         except ImportError:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
-        return acb_theta_jets(z, tau, ord, square=square)
+        return acb_theta_jets(z, tau, ord)
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index f615af89..cf3696db 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -97,26 +97,48 @@ def acb_theta(acb_mat z, acb_mat tau, ulong square=False):
     return acb_mat([res])
 
 
-def acb_theta_jets(acb_mat z, acb_mat tau, slong ord, ulong square=False):
+def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     r"""
-    Corrected wrapper for acb_theta_jet to handle multivariate Taylor expansions.
+    Computes the coefficients of the Taylor expansion of the vector valued Riemann 
+    theta function `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares.
+
+    This is a wrapper for the C-function
+    `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
+    and it follows the same conventions for the ordering of the theta characteristics.
+
+    This should be used via the method :meth:`.acb_mat.theta_jets`, explicitly ``tau.theta_jets(z, ord)``.
+
+        >>> from flint import acb, acb_mat, showgood, ctx
+        >>> z = acb(1+1j); tau = acb(1.25+3j)
+        >>> acb_mat([[tau]]).theta_jets(acb_mat([[z]]),2)  # doctest: +SKIP
+        [[0.969443038779670 +/- 5.67e-16] + [-0.0305569612081680 +/- 5.13e-17]j, 
+        [-0.191993710594950 +/- 4.89e-16] + [0.191993710747776 +/- 7.42e-16]j,   
+        [0.60317023860834 +/- 2.93e-15] + [0.60317023764810 +/- 4.86e-15]j]
+        [[1.03055696119601 +/- 3.89e-15] + [0.0305569612081680 +/- 5.13e-17]j, 
+        [0.191993710594950 +/- 4.89e-16] + [-0.191993710442123 +/- 4.62e-16]j, 
+        [-0.60317023668787 +/- 5.90e-15] + [-0.60317023764810 +/- 4.86e-15]j]
+        [[-1.22079026757697 +/- 4.36e-15] + [-1.82705551679115 +/- 5.17e-15]j,   
+        [-5.71849316258739 +/- 7.02e-15] + [3.82088827346268 +/- 5.75e-15]j,     
+        [6.0241074288587 +/- 3.88e-14] + [9.0163253443780 +/- 2.05e-14]j]
+        [[-1.82023591012499 +/- 2.67e-15] + [1.21625195015448 +/- 4.14e-15]j,     
+        [3.8353056542516 +/- 4.99e-14] + [5.73981078971270 +/- 6.74e-15]j,    
+        [8.9823364151977 +/- 2.44e-14] + [-6.0022138700195 +/- 3.72e-14]j]
     """
     g = tau.nrows()
     if g == 0:
         return []
 
-    # 1. Calculate the length of the jet for one characteristic
+    # Calculate the length of the jet for one characteristic
     # This is the number of multi-indices (alpha) such that |alpha| < ord
     cdef slong nj = acb_theta_jet_nb(g, ord)
     
-    # 2. Total number of characteristics
+    # Total number of characteristics
     cdef slong nb = 1 << (2 * g)
     
-    # 3. Total number of acb elements to allocate
+    # Total number of acb elements to allocate
     # FLINT stores nj coefficients for each of the nb characteristics
     cdef slong total_size = nb * nj
 
-    # Convert input z to acb_ptr
     cdef acb_ptr zvec = _acb_vec_init(g)
     cdef slong i, j
     for i in range(g):
@@ -126,6 +148,7 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord, ulong square=False):
     cdef slong nb_in = 1  # Number of input z vectors
     cdef ulong ab = 0     # Base characteristic
     cdef ulong all = True # Compute all 2^2g characteristics
+    cdef ulong square = False # Don't compute the squares of the thetas.
 
     # Initialize the output buffer
     cdef acb_ptr theta = _acb_vec_init(total_size)
@@ -134,7 +157,7 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord, ulong square=False):
     # Note: Computes all partial derivatives up to total order 'ord'
     acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
 
-    # 4. Copy the output into a structured format
+    # Copy the output into a structured format
     # We return a list of lists: res[char_idx] = [coeff_0, coeff_1, ...]
     res = []
     cdef acb r
@@ -151,50 +174,4 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord, ulong square=False):
     _acb_vec_clear(zvec, g)
     _acb_vec_clear(theta, total_size)
 
-    return res
-
-
-#def acb_theta_jet(acb_mat z, acb_mat tau, slong ord, ulong square=False):
-    #    r"""
-    #    Computes derivatives of the vector valued Riemann theta function
-    #    `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares.
-    #
-    #    This is a wrapper for the C-function
-    #    `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
-    #    and it follows the same conventions for the ordering of the theta characteristics.
-    #
-    #    This should be used via the method :meth:`.acb_mat.theta`, explicitly ``tau.theta(z)``.
-    #    """
-    #    g = tau.nrows()
-    #    assert tau.ncols() == g
-    #    assert z.nrows() == g
-    #    assert z.ncols() == 1
-    #
-    #    # convert input
-    #    cdef acb_ptr zvec
-    #    zvec = _acb_vec_init(g)
-    #    cdef long i
-    #    for i in range(g):
-    #        acb_set(zvec + i, acb_mat_entry(z.val, i, 0))
-    #    cdef slong nb_in = 1 
-    #    cdef ulong ab = 0 
-    #    cdef ulong all = True
-    #
-    #    # initialize the output
-    #    cdef slong nb = 1 << (2 * g)
-    #    # 1. Calculate the length of the jet for one characteristic
-    #    # This is the number of multi-indices (alpha) such that |alpha| < ord
-    #    cdef slong nj = acb_theta_jet_len(g, ord)
-    #    cdef acb_ptr theta = _acb_vec_init(nb)
-    #
-    #    acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
-    #    _acb_vec_clear(zvec, g)
-    #    # copy the output
-    #    res = []
-    #    cdef acb r
-    #    for i in range(nb):
-    #        r = acb.__new__(acb)
-    #        acb_set(r.val, theta + i)
-    #        res.append(r)
-    #    _acb_vec_clear(theta, nb)
-    #    return acb_mat([res])
+    return acb_mat(res)
\ No newline at end of file

From 2c390d8567f596f85bcf42532abd1353e4ef5834 Mon Sep 17 00:00:00 2001
From: Sam Schiavone <sam.schiavone@gmail.com>
Date: Wed, 20 May 2026 10:31:05 -0400
Subject: [PATCH 05/15] fixed linter issues

---
 src/flint/types/acb_theta.pyx | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index cf3696db..a70fefcb 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -99,7 +99,7 @@ def acb_theta(acb_mat z, acb_mat tau, ulong square=False):
 
 def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     r"""
-    Computes the coefficients of the Taylor expansion of the vector valued Riemann 
+    Computes the coefficients of the Taylor expansion of the vector valued Riemann
     theta function `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares.
 
     This is a wrapper for the C-function
@@ -111,17 +111,17 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
         >>> from flint import acb, acb_mat, showgood, ctx
         >>> z = acb(1+1j); tau = acb(1.25+3j)
         >>> acb_mat([[tau]]).theta_jets(acb_mat([[z]]),2)  # doctest: +SKIP
-        [[0.969443038779670 +/- 5.67e-16] + [-0.0305569612081680 +/- 5.13e-17]j, 
-        [-0.191993710594950 +/- 4.89e-16] + [0.191993710747776 +/- 7.42e-16]j,   
+        [[0.969443038779670 +/- 5.67e-16] + [-0.0305569612081680 +/- 5.13e-17]j,
+        [-0.191993710594950 +/- 4.89e-16] + [0.191993710747776 +/- 7.42e-16]j,
         [0.60317023860834 +/- 2.93e-15] + [0.60317023764810 +/- 4.86e-15]j]
-        [[1.03055696119601 +/- 3.89e-15] + [0.0305569612081680 +/- 5.13e-17]j, 
-        [0.191993710594950 +/- 4.89e-16] + [-0.191993710442123 +/- 4.62e-16]j, 
+        [[1.03055696119601 +/- 3.89e-15] + [0.0305569612081680 +/- 5.13e-17]j,
+        [0.191993710594950 +/- 4.89e-16] + [-0.191993710442123 +/- 4.62e-16]j,
         [-0.60317023668787 +/- 5.90e-15] + [-0.60317023764810 +/- 4.86e-15]j]
-        [[-1.22079026757697 +/- 4.36e-15] + [-1.82705551679115 +/- 5.17e-15]j,   
-        [-5.71849316258739 +/- 7.02e-15] + [3.82088827346268 +/- 5.75e-15]j,     
+        [[-1.22079026757697 +/- 4.36e-15] + [-1.82705551679115 +/- 5.17e-15]j,
+        [-5.71849316258739 +/- 7.02e-15] + [3.82088827346268 +/- 5.75e-15]j,
         [6.0241074288587 +/- 3.88e-14] + [9.0163253443780 +/- 2.05e-14]j]
-        [[-1.82023591012499 +/- 2.67e-15] + [1.21625195015448 +/- 4.14e-15]j,     
-        [3.8353056542516 +/- 4.99e-14] + [5.73981078971270 +/- 6.74e-15]j,    
+        [[-1.82023591012499 +/- 2.67e-15] + [1.21625195015448 +/- 4.14e-15]j,
+        [3.8353056542516 +/- 4.99e-14] + [5.73981078971270 +/- 6.74e-15]j,
         [8.9823364151977 +/- 2.44e-14] + [-6.0022138700195 +/- 3.72e-14]j]
     """
     g = tau.nrows()
@@ -131,10 +131,10 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     # Calculate the length of the jet for one characteristic
     # This is the number of multi-indices (alpha) such that |alpha| < ord
     cdef slong nj = acb_theta_jet_nb(g, ord)
-    
+
     # Total number of characteristics
     cdef slong nb = 1 << (2 * g)
-    
+
     # Total number of acb elements to allocate
     # FLINT stores nj coefficients for each of the nb characteristics
     cdef slong total_size = nb * nj
@@ -147,8 +147,8 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     # Parameters for characteristics
     cdef slong nb_in = 1  # Number of input z vectors
     cdef ulong ab = 0     # Base characteristic
-    cdef ulong all = True # Compute all 2^2g characteristics
-    cdef ulong square = False # Don't compute the squares of the thetas.
+    cdef ulong all = True  # Compute all 2^2g characteristics
+    cdef ulong square = False  # Don't compute the squares of the thetas.
 
     # Initialize the output buffer
     cdef acb_ptr theta = _acb_vec_init(total_size)
@@ -174,4 +174,5 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     _acb_vec_clear(zvec, g)
     _acb_vec_clear(theta, total_size)
 
-    return acb_mat(res)
\ No newline at end of file
+    return acb_mat(res)
+

From 3585383e6b8ab139cf5dc80c786064f387c5ae10 Mon Sep 17 00:00:00 2001
From: JHanselman <hanselmanj@hotmail.com>
Date: Wed, 3 Jun 2026 16:54:23 +0200
Subject: [PATCH 06/15] Change matrix creation. Fix (hopefully) linter issues

---
 src/flint/test/test_acb_theta.py | 28 +++++++++++++---------------
 src/flint/types/acb_theta.pyx    | 17 +++++------------
 2 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index c1470d8d..d005b1b9 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -47,20 +47,18 @@ def test_acb_theta_shape_assertions() -> None:
     assert raises(lambda: acb_theta(z, object()), TypeError)  # type: ignore[arg-type]
 
 def test_acb_theta_jets_basic() -> None:
-  if not _has_acb_theta():
-      return
-
-  from flint.types.acb_theta import acb_theta_jets
-
-  z = acb(1 + 1j)
-  tau = acb(1.25 + 3j)
-  zmat = acb_mat([[z]])
-  taumat = acb_mat([[tau]])
-  ord = 2
+    if not _has_acb_theta():
+        return
 
-  direct = acb_theta_jets(zmat, taumat, ord)
-  via_method = taumat.theta_jets(zmat, ord)
-  assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12)
-  assert direct.nrows() == 4
-  assert direct.ncols() == 3
+    from flint.types.acb_theta import acb_theta_jets
+    z = acb(1 + 1j)
+    tau = acb(1.25 + 3j)
+    zmat = acb_mat([[z]])
+    taumat = acb_mat([[tau]])
+    ord = 2
 
+    direct = acb_theta_jets(zmat, taumat, ord)
+    via_method = taumat.theta_jets(zmat, ord)
+    assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12)
+    assert direct.nrows() == 4
+    assert direct.ncols() == 3
\ No newline at end of file
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index cf3696db..83f2422b 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -126,7 +126,7 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     """
     g = tau.nrows()
     if g == 0:
-        return []
+        return acb_mat(0,0)
 
     # Calculate the length of the jet for one characteristic
     # This is the number of multi-indices (alpha) such that |alpha| < ord
@@ -157,21 +157,14 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     # Note: Computes all partial derivatives up to total order 'ord'
     acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
 
-    # Copy the output into a structured format
-    # We return a list of lists: res[char_idx] = [coeff_0, coeff_1, ...]
-    res = []
-    cdef acb r
+   # Copy the output into a structured format
+    res_mat = acb_mat(nb, nj)
     for i in range(nb):
-        char_jet = []
         for j in range(nj):
-            r = acb.__new__(acb)
-            # Offset: characteristic index * length of one jet + coefficient index
-            acb_set(r.val, theta + (i * nj + j))
-            char_jet.append(r)
-        res.append(char_jet)
+            acb_set(acb_mat_entry(res_mat.val, i, j), theta + (i * nj + j)) 
 
     # Cleanup
     _acb_vec_clear(zvec, g)
     _acb_vec_clear(theta, total_size)
 
-    return acb_mat(res)
\ No newline at end of file
+    return res_mat
\ No newline at end of file

From 48ee5c709585d37d6a3e34b13eed448038a74e67 Mon Sep 17 00:00:00 2001
From: JHanselman <hanselmanj@hotmail.com>
Date: Wed, 3 Jun 2026 16:58:54 +0200
Subject: [PATCH 07/15] Empty lines

---
 src/flint/types/acb_theta.pyx | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index 6d19df4a..a1ca0d44 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -167,5 +167,4 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     _acb_vec_clear(zvec, g)
     _acb_vec_clear(theta, total_size)
 
-    return res_mat
-
+    return res_mat
\ No newline at end of file

From 6bb8c52e604c0288bf83571e4ec50724a70a5cdc Mon Sep 17 00:00:00 2001
From: Sam Schiavone <sam.schiavone@gmail.com>
Date: Wed, 10 Jun 2026 10:31:03 -0400
Subject: [PATCH 08/15] fix linter; add function to doc generation file

---
 doc/source/acb_theta.rst      | 2 ++
 src/flint/types/acb_theta.pyx | 9 +++++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/doc/source/acb_theta.rst b/doc/source/acb_theta.rst
index f9ccfdd2..b89cd759 100644
--- a/doc/source/acb_theta.rst
+++ b/doc/source/acb_theta.rst
@@ -3,3 +3,5 @@
 
 .. autofunction :: flint.types.acb_theta.acb_theta
 
+.. autofunction :: flint.types.acb_theta.acb_theta_jets
+
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index a1ca0d44..3cea1a10 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -126,7 +126,7 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     """
     g = tau.nrows()
     if g == 0:
-        return acb_mat(0,0)
+        return acb_mat(0, 0)
 
     # Calculate the length of the jet for one characteristic
     # This is the number of multi-indices (alpha) such that |alpha| < ord
@@ -157,14 +157,15 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     # Note: Computes all partial derivatives up to total order 'ord'
     acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
 
-   # Copy the output into a structured format
+    # Copy the output into a structured format
     res_mat = acb_mat(nb, nj)
     for i in range(nb):
         for j in range(nj):
-            acb_set(acb_mat_entry(res_mat.val, i, j), theta + (i * nj + j)) 
+            acb_set(acb_mat_entry(res_mat.val, i, j), theta + (i * nj + j))
 
     # Cleanup
     _acb_vec_clear(zvec, g)
     _acb_vec_clear(theta, total_size)
 
-    return res_mat
\ No newline at end of file
+    return res_mat
+

From e910da52266e7b574f8c7df926cbec698f0a7b3c Mon Sep 17 00:00:00 2001
From: Sam Schiavone <sam.schiavone@gmail.com>
Date: Wed, 17 Jun 2026 10:43:03 -0400
Subject: [PATCH 09/15] add _has_acb_theta_jet and fix test

---
 README.md                        | 3 +++
 src/flint/__init__.py            | 3 +++
 src/flint/test/test_acb_theta.py | 4 ++--
 src/flint/types/acb_theta.pyx    | 1 -
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index fd9f22ec..4755f041 100644
--- a/README.md
+++ b/README.md
@@ -173,6 +173,8 @@ Contributors (0.9.0):
 - Oscar Benjamin (OB)
 - Daniel Simmons-Marengo (DSM)
 - ForeverHaibara (FH)
+- Jeroen Hanselman (JH)
+- Sam Schiavone (SS)
 
 Changes (0.9.0):
 
@@ -198,6 +200,7 @@ Changes (0.9.0):
 - [gh-359](https://github.com/flintlib/python-flint/pull/359),
   Sort factorisations of all mpoly types. (OB)
 - [gh-374](https://github.com/flintlib/python-flint/pull/374), Fixed a bug in `nmod.__hash__`. (FH)
+- [gh-392](https://github.com/flintlib/python-flint/pull/392), Add interface to `acb_theta_jet` (JH, SS)
 
 0.8.0
 -----
diff --git a/src/flint/__init__.py b/src/flint/__init__.py
index e7811272..391ffa4d 100644
--- a/src/flint/__init__.py
+++ b/src/flint/__init__.py
@@ -66,6 +66,9 @@ def _flint_version_at_least(major: int, minor: int) -> bool:
 def _has_acb_theta() -> bool:
     return _flint_version_at_least(3, 1)
 
+def _has_acb_theta_jet() -> bool:
+    return _flint_version_at_least(3, 3)
+
 
 __all__ = [
     "ctx",
diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index d005b1b9..9c1aacb5 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -47,7 +47,7 @@ def test_acb_theta_shape_assertions() -> None:
     assert raises(lambda: acb_theta(z, object()), TypeError)  # type: ignore[arg-type]
 
 def test_acb_theta_jets_basic() -> None:
-    if not _has_acb_theta():
+    if not _has_acb_theta_jet():
         return
 
     from flint.types.acb_theta import acb_theta_jets
@@ -61,4 +61,4 @@ def test_acb_theta_jets_basic() -> None:
     via_method = taumat.theta_jets(zmat, ord)
     assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12)
     assert direct.nrows() == 4
-    assert direct.ncols() == 3
\ No newline at end of file
+    assert direct.ncols() == 3
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index 3cea1a10..fa4173fb 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -168,4 +168,3 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     _acb_vec_clear(theta, total_size)
 
     return res_mat
-

From 4d71301cce3978cd515358c51121576d2d981820 Mon Sep 17 00:00:00 2001
From: Sam Schiavone <sschiavo@mit.edu>
Date: Wed, 24 Jun 2026 10:33:12 -0400
Subject: [PATCH 10/15] add missing import of _has_acb_theta_jet

---
 src/flint/test/test_acb_theta.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index 9c1aacb5..c80b1247 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from flint import _has_acb_theta, acb, acb_mat
+from flint import _has_acb_theta, _has_acb_theta_jet acb, acb_mat
 from flint.test.helpers import is_close_acb_mat as is_close, raises
 
 

From bf02f5203c10f20eda489ca625a3db980bbe1d55 Mon Sep 17 00:00:00 2001
From: Oscar Benjamin <oscar.j.benjamin@gmail.com>
Date: Sat, 27 Jun 2026 16:46:39 +0100
Subject: [PATCH 11/15] acb_theta_jet: fix syntax error and annotations

---
 src/flint/test/test_acb_theta.py | 3 ++-
 src/flint/types/acb_mat.pyi      | 2 +-
 src/flint/types/acb_mat.pyx      | 2 ++
 src/flint/types/acb_theta.pyi    | 2 +-
 src/flint/types/acb_theta.pyx    | 2 ++
 5 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index c80b1247..92a39f8c 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from flint import _has_acb_theta, _has_acb_theta_jet acb, acb_mat
+from flint import _has_acb_theta, _has_acb_theta_jet, acb, acb_mat
 from flint.test.helpers import is_close_acb_mat as is_close, raises
 
 
@@ -46,6 +46,7 @@ def test_acb_theta_shape_assertions() -> None:
     assert raises(lambda: acb_theta(object(), tau), TypeError)  # type: ignore[arg-type]
     assert raises(lambda: acb_theta(z, object()), TypeError)  # type: ignore[arg-type]
 
+
 def test_acb_theta_jets_basic() -> None:
     if not _has_acb_theta_jet():
         return
diff --git a/src/flint/types/acb_mat.pyi b/src/flint/types/acb_mat.pyi
index 755071a4..a0f23686 100644
--- a/src/flint/types/acb_mat.pyi
+++ b/src/flint/types/acb_mat.pyi
@@ -97,7 +97,7 @@ class acb_mat(flint_mat[acb]):
 
     def theta(self, z: acb_mat, square: bool = False) -> acb_mat: ...
 
-    def theta_jets(self, z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
+    def theta_jets(self, z: acb_mat, ord: int) -> acb_mat: ...
 
 
     def str(self, n: int = 0, radius: bool = True, more: bool = False, condense: int = 0) -> _str: ...
diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index 0b67f2a4..a4e81989 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -851,6 +851,8 @@ cdef class acb_mat(flint_mat):
         This follows the same conventions of the C-function
         `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
         for the ordering of the theta characteristics.
+        The result is an ``acb_mat`` with one row for each theta characteristic and
+        one column for each jet coefficient.
 
         """
         try:
diff --git a/src/flint/types/acb_theta.pyi b/src/flint/types/acb_theta.pyi
index 5d06b495..f921720e 100644
--- a/src/flint/types/acb_theta.pyi
+++ b/src/flint/types/acb_theta.pyi
@@ -5,4 +5,4 @@ from flint.types.acb_mat import acb_mat
 
 def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
 
-def acb_theta_jets(z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ...
+def acb_theta_jets(z: acb_mat, tau: acb_mat, ord: int) -> acb_mat: ...
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index fa4173fb..66ad5b44 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -105,6 +105,8 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     This is a wrapper for the C-function
     `acb_theta_jet <https://flintlib.org/doc/acb_theta.html#c.acb_theta_jet>`_
     and it follows the same conventions for the ordering of the theta characteristics.
+    The result is an ``acb_mat`` with one row for each theta characteristic and
+    one column for each jet coefficient.
 
     This should be used via the method :meth:`.acb_mat.theta_jets`, explicitly ``tau.theta_jets(z, ord)``.
 

From 7c68c594c9d16e3d3571292c86d4b9ed4e423740 Mon Sep 17 00:00:00 2001
From: Oscar Benjamin <oscar.j.benjamin@gmail.com>
Date: Sat, 27 Jun 2026 17:24:44 +0100
Subject: [PATCH 12/15] acb_theta_jet: simplify conditional existence

---
 src/flint/__init__.py                      | 19 ------------------
 src/flint/flintlib/functions/acb_theta.pxd | 23 +++++++++++++++++++---
 src/flint/test/test_acb_mat.py             | 18 ++++++++---------
 src/flint/test/test_acb_theta.py           | 14 +++++--------
 src/flint/test/test_docstrings.py          |  2 +-
 src/flint/types/acb_mat.pyx                |  6 +++++-
 src/flint/types/acb_theta.pyi              | 14 +++++++++++--
 src/flint/types/acb_theta.pyx              | 12 +++++++++--
 8 files changed, 62 insertions(+), 46 deletions(-)

diff --git a/src/flint/__init__.py b/src/flint/__init__.py
index 391ffa4d..87155e8f 100644
--- a/src/flint/__init__.py
+++ b/src/flint/__init__.py
@@ -51,25 +51,6 @@
 __version__ = "0.8.0"
 
 
-def _flint_version_at_least(major: int, minor: int) -> bool:
-    version_parts = __FLINT_VERSION__.split(".")
-    if len(version_parts) < 2:
-        return False
-    try:
-        current_major = int(version_parts[0])
-        current_minor = int(version_parts[1])
-    except ValueError:
-        return False
-    return (current_major, current_minor) >= (major, minor)
-
-
-def _has_acb_theta() -> bool:
-    return _flint_version_at_least(3, 1)
-
-def _has_acb_theta_jet() -> bool:
-    return _flint_version_at_least(3, 3)
-
-
 __all__ = [
     "ctx",
     "fmpz",
diff --git a/src/flint/flintlib/functions/acb_theta.pxd b/src/flint/flintlib/functions/acb_theta.pxd
index fa4a6713..7f6c118d 100644
--- a/src/flint/flintlib/functions/acb_theta.pxd
+++ b/src/flint/flintlib/functions/acb_theta.pxd
@@ -10,12 +10,29 @@ from flint.flintlib.types.fmpz cimport fmpz_mat_struct, fmpz_mat_t
 # unknown type acb_theta_ctx_z_t
 # unknown type acb_theta_sum_worker_t
 
+cdef extern from *:
+    """
+    #include "flint/flint.h"
+    #include "flint/acb_theta.h"
+
+    #if __FLINT_RELEASE >= 30100 /* Flint 3.1.0 or later */
+    #define compat_acb_theta_all(...) acb_theta_all(__VA_ARGS__)
+    #else
+    #define compat_acb_theta_all(...) (void)0
+    #endif
+
+    #if __FLINT_RELEASE >= 30300 /* Flint 3.3.0 or later */
+    #define compat_acb_theta_jet(...) acb_theta_jet(__VA_ARGS__)
+    #else
+    #define compat_acb_theta_jet(...) (void)0
+    #endif
+    """
+    void compat_acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
+    void compat_acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong ord, ulong ab, int all, int sqr, slong prec)
+
 
 cdef extern from "flint/acb_theta.h":
     ulong acb_theta_char_set_slong_vec(const slong * vec, slong len)
-    void acb_theta_one(acb_t th, acb_srcptr z, const acb_mat_t tau, ulong ab, slong prec)
-    void acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
-    void acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong ord, ulong ab, int all, int sqr, slong prec)
     slong sp2gz_dim(const fmpz_mat_t mat)
     void sp2gz_set_blocks(fmpz_mat_t mat, const fmpz_mat_t alpha, const fmpz_mat_t beta, const fmpz_mat_t gamma, const fmpz_mat_t delta)
     void sp2gz_j(fmpz_mat_t mat)
diff --git a/src/flint/test/test_acb_mat.py b/src/flint/test/test_acb_mat.py
index d13ab095..9751006e 100644
--- a/src/flint/test/test_acb_mat.py
+++ b/src/flint/test/test_acb_mat.py
@@ -1,10 +1,13 @@
 from __future__ import annotations
 
-import sys
-from unittest.mock import patch
-
-from flint import _has_acb_theta, acb, acb_mat, arb_mat, ctx, fmpq_mat, fmpz_mat
-from flint.test.helpers import is_close_acb, is_close_acb_mat as is_close, is_close_arb_mat, raises
+from flint import acb, acb_mat, arb_mat, ctx, fmpq_mat, fmpz_mat
+from flint.test.helpers import (
+    is_close_acb,
+    is_close_acb_mat as is_close,
+    is_close_arb_mat,
+    raises,
+)
+from flint.types.acb_theta import acb_theta
 
 
 class _DummyMatrix:
@@ -304,7 +307,7 @@ def test_acb_mat_eig_theta_and_helper() -> None:
     assert acb_mat(0, 0).eig() == []
 
     tau = acb_mat([[1j]])
-    if _has_acb_theta():
+    if acb_theta is not None:
         theta_vals = acb_mat.theta(tau, acb_mat([[0]]))
         assert isinstance(theta_vals, acb_mat)
         assert theta_vals.nrows() == 1
@@ -312,9 +315,6 @@ def test_acb_mat_eig_theta_and_helper() -> None:
     else:
         assert raises(lambda: acb_mat.theta(tau, acb_mat([[0]])), NotImplementedError)
 
-    with patch.dict(sys.modules, {"flint.types.acb_theta": None}):
-        assert raises(lambda: acb_mat.theta(tau, acb_mat([[0]])), NotImplementedError)
-
     assert is_close(a, [[1, 0], [0, 2]]) is True
     assert is_close(a, acb_mat([[1, 0], [0, 2]])) is True
     assert is_close(a, [[1, 0, 0], [0, 2, 0]]) is False
diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index 92a39f8c..81c9643f 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from flint import _has_acb_theta, _has_acb_theta_jet, acb, acb_mat
+from flint import acb, acb_mat
 from flint.test.helpers import is_close_acb_mat as is_close, raises
+from flint.types.acb_theta import acb_theta, acb_theta_jets
 
 
 def test_acb_theta_basic() -> None:
-    if not _has_acb_theta():
+    if acb_theta is None:
         return
 
-    from flint.types.acb_theta import acb_theta
-
     z = acb(1 + 1j)
     tau = acb(1.25 + 3j)
     zmat = acb_mat([[z]])
@@ -27,11 +26,9 @@ def test_acb_theta_basic() -> None:
 
 
 def test_acb_theta_shape_assertions() -> None:
-    if not _has_acb_theta():
+    if acb_theta is None:
         return
 
-    from flint.types.acb_theta import acb_theta
-
     z = acb_mat([[0]])
     tau_bad = acb_mat([[1j, 0]])
     assert raises(lambda: acb_theta(z, tau_bad), AssertionError)
@@ -48,10 +45,9 @@ def test_acb_theta_shape_assertions() -> None:
 
 
 def test_acb_theta_jets_basic() -> None:
-    if not _has_acb_theta_jet():
+    if acb_theta_jets is None:
         return
 
-    from flint.types.acb_theta import acb_theta_jets
     z = acb(1 + 1j)
     tau = acb(1.25 + 3j)
     zmat = acb_mat([[z]])
diff --git a/src/flint/test/test_docstrings.py b/src/flint/test/test_docstrings.py
index 66de7210..8c2b8c56 100644
--- a/src/flint/test/test_docstrings.py
+++ b/src/flint/test/test_docstrings.py
@@ -10,7 +10,7 @@
 test_flint_at_least = {
     "flint.types._gr.gr_ctx.gens": 30100,
     "flint.types._gr.gr_ctx.neg": 30100,
-    "flint.types.acb_theta.acb_theta": 30300,
+    "flint.types.acb_theta.acb_theta": 30100,
 }
 
 
diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index a4e81989..d0f891ea 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -838,6 +838,8 @@ cdef class acb_mat(flint_mat):
             from .acb_theta import acb_theta
         except ImportError:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
+        if acb_theta is None:
+            raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
         return acb_theta(z, tau, square=square)
 
     def theta_jets(tau, z, ord):
@@ -858,5 +860,7 @@ cdef class acb_mat(flint_mat):
         try:
             from .acb_theta import acb_theta_jets
         except ImportError:
-            raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
+            raise NotImplementedError("acb_mat.theta_jets needs Flint >= 3.3.0")
+        if acb_theta_jets is None:
+            raise NotImplementedError("acb_mat.theta_jets needs Flint >= 3.3.0")
         return acb_theta_jets(z, tau, ord)
diff --git a/src/flint/types/acb_theta.pyi b/src/flint/types/acb_theta.pyi
index f921720e..e5bc0239 100644
--- a/src/flint/types/acb_theta.pyi
+++ b/src/flint/types/acb_theta.pyi
@@ -1,8 +1,18 @@
 from __future__ import annotations
 
+from typing import Protocol
+
 from flint.types.acb_mat import acb_mat
 
 
-def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
+class _AcbThetaFunc(Protocol):
+    def __call__(self, z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
+
+
+class _AcbThetaJetsFunc(Protocol):
+    def __call__(self, z: acb_mat, tau: acb_mat, ord: int) -> acb_mat: ...
+
+
+acb_theta: _AcbThetaFunc | None
 
-def acb_theta_jets(z: acb_mat, tau: acb_mat, ord: int) -> acb_mat: ...
+acb_theta_jets: _AcbThetaJetsFunc | None
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index 66ad5b44..2b3d36ea 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -1,4 +1,5 @@
 from flint.flint_base.flint_context cimport getprec
+from flint.flint_base.flint_base import FLINT_RELEASE
 from flint.types.acb cimport acb
 from flint.types.acb_mat cimport acb_mat
 from flint.flintlib.functions.acb cimport *
@@ -84,7 +85,7 @@ def acb_theta(acb_mat z, acb_mat tau, ulong square=False):
     cdef slong nb = 1 << (2 * g)
     cdef acb_ptr theta = _acb_vec_init(nb)
 
-    acb_theta_all(theta, zvec, tau.val, square, getprec())
+    compat_acb_theta_all(theta, zvec, tau.val, square, getprec())
     _acb_vec_clear(zvec, g)
     # copy the output
     res = []
@@ -157,7 +158,7 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
 
     # Call the FLINT C function
     # Note: Computes all partial derivatives up to total order 'ord'
-    acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
+    compat_acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec())
 
     # Copy the output into a structured format
     res_mat = acb_mat(nb, nj)
@@ -170,3 +171,10 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
     _acb_vec_clear(theta, total_size)
 
     return res_mat
+
+
+if FLINT_RELEASE < 30100:
+    acb_theta = None
+
+if FLINT_RELEASE < 30300:
+    acb_theta_jets = None

From 44f5847904e85c283ce4525783f84c39e4143f2c Mon Sep 17 00:00:00 2001
From: Oscar Benjamin <oscar.j.benjamin@gmail.com>
Date: Sat, 27 Jun 2026 17:31:43 +0100
Subject: [PATCH 13/15] acb_theta: don't use delayerd imports

---
 src/flint/types/acb_mat.pyi |  2 --
 src/flint/types/acb_mat.pyx | 13 +++----------
 2 files changed, 3 insertions(+), 12 deletions(-)

diff --git a/src/flint/types/acb_mat.pyi b/src/flint/types/acb_mat.pyi
index a0f23686..75a4e523 100644
--- a/src/flint/types/acb_mat.pyi
+++ b/src/flint/types/acb_mat.pyi
@@ -96,10 +96,8 @@ class acb_mat(flint_mat[acb]):
     ) -> Any: ...
 
     def theta(self, z: acb_mat, square: bool = False) -> acb_mat: ...
-
     def theta_jets(self, z: acb_mat, ord: int) -> acb_mat: ...
 
-
     def str(self, n: int = 0, radius: bool = True, more: bool = False, condense: int = 0) -> _str: ...
     def repr(self) -> _str: ...
     def __str__(self) -> _str: ...
diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx
index d0f891ea..dec64eb7 100644
--- a/src/flint/types/acb_mat.pyx
+++ b/src/flint/types/acb_mat.pyx
@@ -5,9 +5,8 @@ from flint.types.arb_mat cimport arb_mat
 from flint.types.fmpz_mat cimport fmpz_mat
 from flint.types.fmpq_mat cimport fmpq_mat
 from flint.types.arb cimport arb
-from flint.types.acb cimport acb
+from flint.types.acb cimport acb, any_as_acb
 from flint.types.acb_poly cimport acb_poly
-from flint.types.acb cimport any_as_acb
 from flint.types.fmpz cimport fmpz
 from flint.types.fmpq cimport fmpq
 
@@ -38,6 +37,8 @@ from flint.flintlib.types.acb cimport (
 
 cimport cython
 
+from flint.types.acb_theta import acb_theta, acb_theta_jets
+
 cdef acb_mat_coerce_operands(x, y):
     if isinstance(y, (fmpz_mat, fmpq_mat, arb_mat)):
         return x, acb_mat(y)
@@ -834,10 +835,6 @@ cdef class acb_mat(flint_mat):
         for the ordering of the theta characteristics.
 
         """
-        try:
-            from .acb_theta import acb_theta
-        except ImportError:
-            raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
         if acb_theta is None:
             raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0")
         return acb_theta(z, tau, square=square)
@@ -857,10 +854,6 @@ cdef class acb_mat(flint_mat):
         one column for each jet coefficient.
 
         """
-        try:
-            from .acb_theta import acb_theta_jets
-        except ImportError:
-            raise NotImplementedError("acb_mat.theta_jets needs Flint >= 3.3.0")
         if acb_theta_jets is None:
             raise NotImplementedError("acb_mat.theta_jets needs Flint >= 3.3.0")
         return acb_theta_jets(z, tau, ord)

From 221d6b7912575fd1f4bca0a4ebf49046049f5f98 Mon Sep 17 00:00:00 2001
From: Oscar Benjamin <oscar.j.benjamin@gmail.com>
Date: Sat, 27 Jun 2026 17:57:45 +0100
Subject: [PATCH 14/15] acb_theta: build the module unconditionally

---
 meson.build                      |  3 --
 setup.py                         |  1 +
 src/flint/test/test_acb_theta.py |  5 +--
 src/flint/types/acb_theta.pyx    | 59 ++++++++++++++++++++++++++++++--
 src/flint/types/meson.build      |  5 +--
 5 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/meson.build b/meson.build
index fbabce65..61b9d036 100644
--- a/meson.build
+++ b/meson.build
@@ -73,9 +73,6 @@ endif
 # flint.pc was missing -lflint until Flint 3.1.0
 if flint_dep.version().version_compare('<3.1')
   flint_dep = cc.find_library('flint')
-  have_acb_theta = false
-else
-  have_acb_theta = true
 endif
 
 pyflint_deps = [gmp_dep, mpfr_dep, flint_dep]
diff --git a/setup.py b/setup.py
index 669bcfe4..c9e5e22d 100644
--- a/setup.py
+++ b/setup.py
@@ -119,6 +119,7 @@
     ("flint.types.acb_poly", ["src/flint/types/acb_poly.pyx"]),
     ("flint.types.acb_mat", ["src/flint/types/acb_mat.pyx"]),
     ("flint.types.acb_series", ["src/flint/types/acb_series.pyx"]),
+    ("flint.types.acb_theta", ["src/flint/types/acb_theta.pyx"]),
 
     ("flint.types.dirichlet", ["src/flint/types/dirichlet.pyx"]),
 
diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py
index 81c9643f..4446002c 100644
--- a/src/flint/test/test_acb_theta.py
+++ b/src/flint/test/test_acb_theta.py
@@ -8,19 +8,20 @@
 def test_acb_theta_basic() -> None:
     if acb_theta is None:
         return
+    theta = acb_theta
 
     z = acb(1 + 1j)
     tau = acb(1.25 + 3j)
     zmat = acb_mat([[z]])
     taumat = acb_mat([[tau]])
 
-    direct = acb_theta(zmat, taumat)
+    direct = theta(zmat, taumat)
     via_method = taumat.theta(zmat)
     assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12)
     assert direct.nrows() == 1
     assert direct.ncols() == 4
 
-    squared = acb_theta(zmat, taumat, square=True)
+    squared = theta(zmat, taumat, square=True)
     assert squared.nrows() == 1
     assert squared.ncols() == 4
 
diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx
index 2b3d36ea..240faaf6 100644
--- a/src/flint/types/acb_theta.pyx
+++ b/src/flint/types/acb_theta.pyx
@@ -4,10 +4,65 @@ from flint.types.acb cimport acb
 from flint.types.acb_mat cimport acb_mat
 from flint.flintlib.functions.acb cimport *
 from flint.flintlib.types.acb cimport (
+    acb_mat_t,
     acb_mat_entry,
+    acb_ptr,
+    acb_srcptr,
 )
 from flint.flintlib.functions.acb_mat cimport *
-from flint.flintlib.functions.acb_theta cimport *
+from flint.flintlib.types.flint cimport slong, ulong
+
+
+cdef extern from *:
+    """
+    #include "flint/flint.h"
+    #include "flint/acb.h"
+    #include "flint/acb_mat.h"
+
+    #if __FLINT_RELEASE >= 30100 /* Flint 3.1.0 or later */
+    #include "flint/acb_theta.h"
+    static inline void
+    compat_acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
+    {
+        acb_theta_all(th, z, tau, sqr, prec);
+    }
+    #else
+    static inline void
+    compat_acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
+    {
+    }
+    #endif
+
+    #if __FLINT_RELEASE >= 30300 /* Flint 3.3.0 or later */
+    static inline slong
+    compat_acb_theta_jet_nb(slong ord, slong g)
+    {
+        return acb_theta_jet_nb(g, ord);
+    }
+
+    static inline void
+    compat_acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau,
+            slong ord, ulong ab, int all, int sqr, slong prec)
+    {
+        acb_theta_jet(th, zs, nb, tau, ord, ab, all, sqr, prec);
+    }
+    #else
+    static inline slong
+    compat_acb_theta_jet_nb(slong ord, slong g)
+    {
+        return 0;
+    }
+
+    static inline void
+    compat_acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau,
+            slong ord, ulong ab, int all, int sqr, slong prec)
+    {
+    }
+    #endif
+    """
+    void compat_acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
+    slong compat_acb_theta_jet_nb(slong ord, slong g)
+    void compat_acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong ord, ulong ab, int all, int sqr, slong prec)
 
 
 def acb_theta(acb_mat z, acb_mat tau, ulong square=False):
@@ -133,7 +188,7 @@ def acb_theta_jets(acb_mat z, acb_mat tau, slong ord):
 
     # Calculate the length of the jet for one characteristic
     # This is the number of multi-indices (alpha) such that |alpha| < ord
-    cdef slong nj = acb_theta_jet_nb(g, ord)
+    cdef slong nj = compat_acb_theta_jet_nb(ord, g)
 
     # Total number of characteristics
     cdef slong nb = 1 << (2 * g)
diff --git a/src/flint/types/meson.build b/src/flint/types/meson.build
index b894ffd9..1e527091 100644
--- a/src/flint/types/meson.build
+++ b/src/flint/types/meson.build
@@ -83,6 +83,7 @@ exts = [
   'arb',
   'arb_poly',
   'arb_mat',
+  'acb_theta',
   'arb_series',
 
   'acb',
@@ -95,10 +96,6 @@ exts = [
   '_gr',
 ]
 
-if have_acb_theta
-  exts += ['acb_theta']
-endif
-
 py.install_sources(
   pyfiles,
   pure: false,

From 2ed8f531f6cbee565bf2b88e58cdf2579c49938e Mon Sep 17 00:00:00 2001
From: Oscar Benjamin <oscar.j.benjamin@gmail.com>
Date: Sat, 27 Jun 2026 18:17:28 +0100
Subject: [PATCH 15/15] Revert changes to auto-generated acb_theta.pxd

---
 src/flint/flintlib/functions/acb_theta.pxd | 23 +++-------------------
 1 file changed, 3 insertions(+), 20 deletions(-)

diff --git a/src/flint/flintlib/functions/acb_theta.pxd b/src/flint/flintlib/functions/acb_theta.pxd
index 7f6c118d..fa4a6713 100644
--- a/src/flint/flintlib/functions/acb_theta.pxd
+++ b/src/flint/flintlib/functions/acb_theta.pxd
@@ -10,29 +10,12 @@ from flint.flintlib.types.fmpz cimport fmpz_mat_struct, fmpz_mat_t
 # unknown type acb_theta_ctx_z_t
 # unknown type acb_theta_sum_worker_t
 
-cdef extern from *:
-    """
-    #include "flint/flint.h"
-    #include "flint/acb_theta.h"
-
-    #if __FLINT_RELEASE >= 30100 /* Flint 3.1.0 or later */
-    #define compat_acb_theta_all(...) acb_theta_all(__VA_ARGS__)
-    #else
-    #define compat_acb_theta_all(...) (void)0
-    #endif
-
-    #if __FLINT_RELEASE >= 30300 /* Flint 3.3.0 or later */
-    #define compat_acb_theta_jet(...) acb_theta_jet(__VA_ARGS__)
-    #else
-    #define compat_acb_theta_jet(...) (void)0
-    #endif
-    """
-    void compat_acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
-    void compat_acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong ord, ulong ab, int all, int sqr, slong prec)
-
 
 cdef extern from "flint/acb_theta.h":
     ulong acb_theta_char_set_slong_vec(const slong * vec, slong len)
+    void acb_theta_one(acb_t th, acb_srcptr z, const acb_mat_t tau, ulong ab, slong prec)
+    void acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
+    void acb_theta_jet(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong ord, ulong ab, int all, int sqr, slong prec)
     slong sp2gz_dim(const fmpz_mat_t mat)
     void sp2gz_set_blocks(fmpz_mat_t mat, const fmpz_mat_t alpha, const fmpz_mat_t beta, const fmpz_mat_t gamma, const fmpz_mat_t delta)
     void sp2gz_j(fmpz_mat_t mat)