diff --git a/src/sparsejac/__init__.py b/src/sparsejac/__init__.py
index a48de17..06a9a07 100644
--- a/src/sparsejac/__init__.py
+++ b/src/sparsejac/__init__.py
@@ -3,5 +3,9 @@
 __version__ = "v0.2.0"
 __author__ = "Martin Schubert <mfschubert@gmail.com>"
 
+from sparsejac.sparsejac import ForwardModeColoring as ForwardModeColoring
+from sparsejac.sparsejac import ReverseModeColoring as ReverseModeColoring
 from sparsejac.sparsejac import jacfwd as jacfwd
+from sparsejac.sparsejac import jacfwd_coloring as jacfwd_coloring
 from sparsejac.sparsejac import jacrev as jacrev
+from sparsejac.sparsejac import jacrev_coloring as jacrev_coloring
diff --git a/src/sparsejac/sparsejac.py b/src/sparsejac/sparsejac.py
index 9f51649..a862fdb 100644
--- a/src/sparsejac/sparsejac.py
+++ b/src/sparsejac/sparsejac.py
@@ -1,6 +1,6 @@
 """Defines functions for efficient computation of sparse Jacobians."""
 
-from typing import Any, Callable, Tuple, Union
+from typing import Any, Callable, NamedTuple, NewType, Tuple, Union
 
 import jax
 import jax.experimental.sparse as jsparse
@@ -14,15 +14,190 @@
 
 ArrayWithOptionalAux = Union[jnp.ndarray, Tuple[jnp.ndarray, Any]]
 
+@jax.tree_util.register_static
+class StaticInt(int):
+    pass
+
+class ReverseModeColoring(NamedTuple):
+    """Precomputed coloring for use with `jacrev`.
+
+    Attributes:
+        sparsity: Sparse matrix identifying nonzero Jacobian locations.
+        output_coloring: Integer array of shape `(nrows,)` assigning each
+            output element (corresponding to a row in the Jacobian) a color
+            index in `[0, ncolors)`.
+        ncolors: Total number of colors used. Static, because it is used as
+            part of an array shape.
+    """
+
+    sparsity: jsparse.BCOO
+    output_coloring: jnp.ndarray
+    ncolors: StaticInt
+
+
+class ForwardModeColoring(NamedTuple):
+    """Precomputed coloring for use with `jacfwd`.
+
+    Attributes:
+        sparsity: Sparse matrix identifying nonzero Jacobian locations.
+        input_coloring: Integer array of shape `(ncols,)` assigning each
+            input element (corresponding to a column in the Jacobian) a color
+            index in `[0, ncolors)`.
+        ncolors: Total number of colors used. Static, because it is used as
+            part of an array shape.
+    """
+
+    sparsity: jsparse.BCOO
+    input_coloring: jnp.ndarray
+    ncolors: StaticInt
+
+
+def jacrev_coloring(
+    sparsity: Union[jsparse.BCOO, ssparse.spmatrix],
+    coloring_strategy: str = _DEFAULT_COLORING_STRATEGY,
+) -> ReverseModeColoring:
+    """Returns the precomputed coloring for use with `jacrev`.
+
+    This function performs the graph-coloring step that identifies structurally
+    independent groups of output elements. Because it uses `scipy` and
+    `networkx`, it is not JAX-traceable and must be called outside of
+    `jax.jit`. The returned `ReverseModeColoring` can then be passed to
+    `jacrev` via the `coloring` keyword argument, making `jacrev` itself
+    callable inside a jitted function without needing
+    `jax.ensure_compile_time_eval`.
+
+    The `output_coloring` field of the returned object encodes a projection
+    matrix `P` of shape `(ncolors, nrows)` where `P[c, i] = 1` iff
+    `output_coloring[i] == c`. This matrix compresses the output before
+    reverse-mode differentiation; the scatter-add
+    `jnp.zeros(ncolors).at[output_coloring].add(y)` is equivalent to `P @ y`
+    and avoids materialising the full 2-D matrix.
+
+    Args:
+        sparsity: Sparse matrix whose specified elements are at locations where
+            the Jacobian is nonzero. Note that the values of `sparsity` are not
+            used. Accepts either a JAX `BCOO` matrix (which must be rank-2 with
+            `n_sparse == 2`) or a SciPy sparse matrix.
+        coloring_strategy: See `networkx.algorithms.coloring.greedy_color`.
+
+    Returns:
+        A `ReverseModeColoring` containing the sparsity, output coloring
+        vector, and number of colors.
+
+    Raises:
+        ValueError: If `sparsity` is not rank-2 or, for `BCOO` inputs,
+            `sparsity.n_sparse != 2`.
+    """
+    if isinstance(sparsity, jsparse.BCOO):
+        if sparsity.ndim != 2:
+            raise ValueError(
+                f"`sparsity` must be rank-2, but got shape of {sparsity.shape}."
+            )
+        if sparsity.n_sparse != 2:
+            raise ValueError(
+                f"`sparsity.n_sparse` must be 2, but got a value of "
+                f"{sparsity.n_sparse}."
+            )
+        sparsity_bcoo = sparsity
+        sparsity_scipy = ssparse.coo_matrix(
+            (sparsity.data, sparsity.indices.T), shape=sparsity.shape
+        )
+    else:
+        if sparsity.ndim != 2:
+            raise ValueError(
+                f"`sparsity` must be rank-2, but got shape of {sparsity.shape}."
+            )
+        sparsity_bcoo = jsparse.BCOO.from_scipy_sparse(sparsity)
+        sparsity_scipy = sparsity
+
+    connectivity = _output_connectivity_from_sparsity(sparsity_scipy)
+    output_coloring_array, ncolors = _greedy_color(connectivity, coloring_strategy)
+    assert output_coloring_array.size == sparsity_bcoo.shape[0]
+    return ReverseModeColoring(
+        sparsity=sparsity_bcoo,
+        output_coloring=jnp.asarray(output_coloring_array),
+        ncolors=StaticInt(ncolors),
+    )
+
+
+def jacfwd_coloring(
+    sparsity: Union[jsparse.BCOO, ssparse.spmatrix],
+    coloring_strategy: str = _DEFAULT_COLORING_STRATEGY,
+) -> ForwardModeColoring:
+    """Returns the precomputed coloring for use with `jacfwd`.
+
+    This function performs the graph-coloring step that identifies structurally
+    independent groups of input elements. Because it uses `scipy` and
+    `networkx`, it is not JAX-traceable and must be called outside of
+    `jax.jit`. The returned `ForwardModeColoring` can then be passed to
+    `jacfwd` via the `coloring` keyword argument, making `jacfwd` itself
+    callable inside a jitted function without needing
+    `jax.ensure_compile_time_eval`.
+
+    The `input_coloring` field of the returned object encodes a basis matrix
+    `B` of shape `(ncols, ncolors)` where `B[j, c] = 1` iff
+    `input_coloring[j] == c`. Column `c` of `B`, used as a JVP tangent vector,
+    is equivalent to `(input_coloring == c).astype(float)`, so the vmap over
+    colors can compute tangents on-the-fly and avoids materialising the full
+    matrix.
+
+    Args:
+        sparsity: Sparse matrix whose specified elements are at locations where
+            the Jacobian is nonzero. Note that the values of `sparsity` are not
+            used. Accepts either a JAX `BCOO` matrix (which must be rank-2 with
+            `n_sparse == 2`) or a SciPy sparse matrix.
+        coloring_strategy: See `networkx.algorithms.coloring.greedy_color`.
+
+    Returns:
+        A `ForwardModeColoring` containing the sparsity, input coloring
+        vector, and number of colors.
+
+    Raises:
+        ValueError: If `sparsity` is not rank-2 or, for `BCOO` inputs,
+            `sparsity.n_sparse != 2`.
+    """
+    if isinstance(sparsity, jsparse.BCOO):
+        if sparsity.ndim != 2:
+            raise ValueError(
+                f"`sparsity` must be rank-2, but got shape of {sparsity.shape}."
+            )
+        if sparsity.n_sparse != 2:
+            raise ValueError(
+                f"`sparsity.n_sparse` must be 2, but got a value of "
+                f"{sparsity.n_sparse}."
+            )
+        sparsity_bcoo = sparsity
+        sparsity_scipy = ssparse.coo_matrix(
+            (sparsity.data, sparsity.indices.T), shape=sparsity.shape
+        )
+    else:
+        if sparsity.ndim != 2:
+            raise ValueError(
+                f"`sparsity` must be rank-2, but got shape of {sparsity.shape}."
+            )
+        sparsity_bcoo = jsparse.BCOO.from_scipy_sparse(sparsity)
+        sparsity_scipy = sparsity
+
+    connectivity = _input_connectivity_from_sparsity(sparsity_scipy)
+    input_coloring_array, ncolors = _greedy_color(connectivity, coloring_strategy)
+    assert input_coloring_array.size == sparsity_bcoo.shape[1]
+    return ForwardModeColoring(
+        sparsity=sparsity_bcoo,
+        input_coloring=jnp.asarray(input_coloring_array),
+        ncolors=StaticInt(ncolors),
+    )
+
 
 def jacrev(
     fn: Callable[[Any], ArrayWithOptionalAux],
-    sparsity: jsparse.BCOO,
+    sparsity: Union[jsparse.BCOO, None] = None,
     argnums: int = 0,
     has_aux: bool = False,
     holomorphic: bool = False,
     allow_int: bool = False,
     coloring_strategy: str = _DEFAULT_COLORING_STRATEGY,
+    *,
+    coloring: Union[ReverseModeColoring, None] = None,
 ) -> Callable[[Any], ArrayWithOptionalAux]:
     """Returns a function which computes the Jacobian of `fn` using reverse mode.
 
@@ -36,6 +211,13 @@ def jacrev(
     project to a lower-dimensional output space, so that reverse-mode
     differentiation can be more efficiently applied.
 
+    Either `sparsity` or `coloring` must be provided, but not both. When
+    `sparsity` is given, the coloring is computed internally via
+    `jacrev_coloring`. When `coloring` is given (a precomputed
+    `ReverseModeColoring`), the graph-coloring step is skipped, which allows
+    this function to be used inside `jax.jit` without
+    `jax.ensure_compile_time_eval`.
+
     Args:
         fn: The function for which the sparse Jacobian is sought. The function can
             have several inputs, one of which is to be differentiated with respect to.
@@ -44,48 +226,44 @@ def jacrev(
             `fn` must be rank-1 with size equal to the row count in `sparsity`.
         sparsity: Sparse matrix whose specified elements are at locations where the
             Jacobian is nonzero. Note that the values of `sparsity` are not used.
+            Mutually exclusive with `coloring`.
         argnums: Specifies the positional argument to differentiate with respect to.
         has_aux: See `jax.jacrev`.
         holomorphic: See `jax.jacrev`.
         allow_int: See `jax.jacrev`.
-        coloring_strategy: See `networkx.algorithms.coloring.greedy_color`.
+        coloring_strategy: See `networkx.algorithms.coloring.greedy_color`. Only
+            used when `sparsity` is provided.
+        coloring: Precomputed coloring from `jacrev_coloring`. Mutually exclusive
+            with `sparsity`.
 
     Returns:
         The function which computes the sparse Jacobian.
+
+    Raises:
+        ValueError: If both or neither of `sparsity` and `coloring` are provided,
+            or if `argnums` is not an integer.
     """
-    if sparsity.ndim != 2:
-        raise ValueError(
-            f"`sparsity` must be rank-2, but got shape of {sparsity.shape}."
-        )
-    if sparsity.n_sparse != 2:
-        raise ValueError(
-            f"`sparsity.n_sparse` must be 2, but got a value of "
-            f"{sparsity.n_sparse}."
-        )
+    if sparsity is not None and coloring is not None:
+        raise ValueError("Provide either `sparsity` or `coloring`, not both.")
+    if sparsity is None and coloring is None:
+        raise ValueError("Exactly one of `sparsity` or `coloring` must be provided.")
     if not isinstance(argnums, int):
         raise ValueError(f"`argnums` must be an integer, but got {argnums}.")
+    if sparsity is not None:
+        coloring = jacrev_coloring(sparsity, coloring_strategy)
+    assert coloring is not None
 
-    # Identify the structurally-independent elements of `fn` output, i.e. obtain
-    # the coloring of the output. Here we must use `scipy` sparse matrices.
-    sparsity_scipy = ssparse.coo_matrix(
-        (sparsity.data, sparsity.indices.T), shape=sparsity.shape
-    )
-    connectivity = _output_connectivity_from_sparsity(sparsity_scipy)
-    output_coloring, ncolors = _greedy_color(connectivity, coloring_strategy)
-    output_coloring = jnp.asarray(output_coloring)
-    assert output_coloring.size == sparsity.shape[0]
-
-    projection_matrix = (
-        jnp.arange(ncolors)[:, jnp.newaxis] == output_coloring[jnp.newaxis, :]
-    )
+    _sparsity = coloring.sparsity
+    output_coloring = coloring.output_coloring
+    ncolors = coloring.ncolors
 
     def jacrev_fn(*args: Any, **kwargs: Any) -> ArrayWithOptionalAux:
         x = args[argnums]
-        if x.shape != (sparsity.shape[1],):
+        if x.shape != (_sparsity.shape[1],):
             raise ValueError(
                 f"`x` must be rank-1 with size matching the number of columns in "
                 f"`sparsity`, but got shape {x.shape} when `sparsity` has shape "
-                f"{sparsity.shape}."
+                f"{_sparsity.shape}."
             )
 
         def _projected_fn(*args: Any, **kwargs: Any) -> ArrayWithOptionalAux:
@@ -93,16 +271,18 @@ def _projected_fn(*args: Any, **kwargs: Any) -> ArrayWithOptionalAux:
                 y, aux = fn(*args, **kwargs)
             else:
                 y = fn(*args, **kwargs)
-            if y.shape != (sparsity.shape[0],):
+            if y.shape != (_sparsity.shape[0],):
                 raise ValueError(
                     f"`fn(x)` must be rank-1 with size matching the number of rows in "
                     f"`sparsity`, but got shape {y.shape} when `sparsity` has shape "
-                    f"{sparsity.shape}."
+                    f"{_sparsity.shape}."
                 )
+            # Equivalent to `P @ y` where P[c, i] = 1 iff output_coloring[i] == c.
+            projected = jnp.zeros(ncolors, dtype=y.dtype).at[output_coloring].add(y)
             if has_aux:
-                return projection_matrix @ y, aux
+                return projected, aux
             else:
-                return projection_matrix @ y
+                return projected
 
         compressed_jac_with_maybe_aux = jax.jacrev(
             _projected_fn,
@@ -114,20 +294,22 @@ def _projected_fn(*args: Any, **kwargs: Any) -> ArrayWithOptionalAux:
 
         if has_aux:
             compressed_jac, aux = compressed_jac_with_maybe_aux
-            return _expand_jacrev_jac(compressed_jac, output_coloring, sparsity), aux
+            return _expand_jacrev_jac(compressed_jac, output_coloring, _sparsity), aux
         else:
             compressed_jac = compressed_jac_with_maybe_aux
-            return _expand_jacrev_jac(compressed_jac, output_coloring, sparsity)
+            return _expand_jacrev_jac(compressed_jac, output_coloring, _sparsity)
 
     return jacrev_fn
 
 
 def jacfwd(
     fn: Callable[[Any], ArrayWithOptionalAux],
-    sparsity: jsparse.BCOO,
+    sparsity: Union[jsparse.BCOO, None] = None,
     argnums: int = 0,
     has_aux: bool = False,
     coloring_strategy: str = _DEFAULT_COLORING_STRATEGY,
+    *,
+    coloring: Union[ForwardModeColoring, None] = None,
 ) -> Callable[[Any], ArrayWithOptionalAux]:
     """Returns a function which computes the Jacobian of `fn` using forward mode.
 
@@ -141,6 +323,13 @@ def jacfwd(
     project to a lower-dimensional input space, so that forward-mode
     differentiation can be more efficiently applied.
 
+    Either `sparsity` or `coloring` must be provided, but not both. When
+    `sparsity` is given, the coloring is computed internally via
+    `jacfwd_coloring`. When `coloring` is given (a precomputed
+    `ForwardModeColoring`), the graph-coloring step is skipped, which allows
+    this function to be used inside `jax.jit` without
+    `jax.ensure_compile_time_eval`.
+
     Args:
         fn: The function for which the sparse Jacobian is sought. The function can
             have several inputs, one of which is to be differentiated with respect to.
@@ -149,45 +338,42 @@ def jacfwd(
             `fn` must be rank-1 with size equal to the row count in `sparsity`.
         sparsity: Sparse matrix whose specified elements are at locations where the
             Jacobian is nonzero. Note that the values of `sparsity` are not used.
+            Mutually exclusive with `coloring`.
         argnums: Specifies the positional argument to differentiate with respect to.
         has_aux: See `jax.jacfwd`.
-        coloring_strategy: See `networkx.algorithms.coloring.greedy_color`.
+        coloring_strategy: See `networkx.algorithms.coloring.greedy_color`. Only
+            used when `sparsity` is provided.
+        coloring: Precomputed coloring from `jacfwd_coloring`. Mutually exclusive
+            with `sparsity`.
 
     Returns:
         The function which computes the sparse Jacobian.
+
+    Raises:
+        ValueError: If both or neither of `sparsity` and `coloring` are provided,
+            or if `argnums` is not an integer.
     """
-    if sparsity.ndim != 2:
-        raise ValueError(
-            f"`sparsity` must be rank-2, but got shape of {sparsity.shape}."
-        )
-    if sparsity.n_sparse != 2:
-        raise ValueError(
-            f"`sparsity.n_sparse` must be 2, but got a value of "
-            f"{sparsity.n_sparse}."
-        )
+    if sparsity is not None and coloring is not None:
+        raise ValueError("Provide either `sparsity` or `coloring`, not both.")
+    if sparsity is None and coloring is None:
+        raise ValueError("Exactly one of `sparsity` or `coloring` must be provided.")
     if not isinstance(argnums, int):
         raise ValueError(f"`argnums` must be an integer, but got {argnums}.")
+    if sparsity is not None:
+        coloring = jacfwd_coloring(sparsity, coloring_strategy)
+    assert coloring is not None
 
-    # Identify the structurally-independent elements of `fn` output, i.e. obtain
-    # the coloring of the output. Here we must use `scipy` sparse matrices.
-    sparsity_scipy = ssparse.coo_matrix(
-        (sparsity.data, sparsity.indices.T), shape=sparsity.shape
-    )
-    connectivity = _input_connectivity_from_sparsity(sparsity_scipy)
-    input_coloring, ncolors = _greedy_color(connectivity, coloring_strategy)
-    input_coloring = jnp.asarray(input_coloring)
-    assert input_coloring.size == sparsity.shape[1]
-
-    basis = jnp.arange(ncolors)[jnp.newaxis, :] == input_coloring[:, jnp.newaxis]
-    basis = basis.astype(float)
+    _sparsity = coloring.sparsity
+    input_coloring = coloring.input_coloring
+    ncolors = coloring.ncolors
 
     def jacfwd_fn(*args: Any, **kwargs: Any) -> ArrayWithOptionalAux:
         x = args[argnums]
-        if x.shape != (sparsity.shape[1],):
+        if x.shape != (_sparsity.shape[1],):
             raise ValueError(
                 f"`x` must be rank-1 with size matching the number of columns in "
                 f"`sparsity`, but got shape {x.shape} when `sparsity` has shape "
-                f"{sparsity.shape}."
+                f"{_sparsity.shape}."
             )
 
         def _fn(x: jnp.ndarray) -> ArrayWithOptionalAux:
@@ -195,36 +381,40 @@ def _fn(x: jnp.ndarray) -> ArrayWithOptionalAux:
             return fn(*args_with_x, **kwargs)
 
         if has_aux:
-
-            def _jvp_fn_with_aux(tangents: jnp.ndarray) -> Tuple[jnp.ndarray, Any]:
+            # Equivalent to vmapping over columns of B where B[j, c] = 1 iff
+            # input_coloring[j] == c.
+            def _jvp_fn_with_aux(c: jnp.ndarray) -> Tuple[jnp.ndarray, Any]:
+                tangents = (input_coloring == c).astype(float)
                 _, tangents_out, aux = jax.jvp(_fn, (x,), (tangents,), has_aux=True)
                 return tangents_out, aux
 
             compressed_jac_transpose, aux = jax.vmap(
-                _jvp_fn_with_aux, in_axes=1, out_axes=(0, None)
-            )(basis)
+                _jvp_fn_with_aux, out_axes=(0, None)
+            )(jnp.arange(ncolors))
 
         else:
-
-            def _jvp_fn(tangents: jnp.ndarray) -> jnp.ndarray:
+            # Equivalent to vmapping over columns of B where B[j, c] = 1 iff
+            # input_coloring[j] == c.
+            def _jvp_fn(c: jnp.ndarray) -> jnp.ndarray:
+                tangents = (input_coloring == c).astype(float)
                 _, tangents_out = jax.jvp(_fn, (x,), (tangents,), has_aux=False)
                 return tangents_out
 
-            compressed_jac_transpose = jax.vmap(_jvp_fn, in_axes=1)(basis)
+            compressed_jac_transpose = jax.vmap(_jvp_fn)(jnp.arange(ncolors))
 
         compressed_jac = compressed_jac_transpose.T
-        if compressed_jac.shape != (sparsity.shape[0], ncolors):
+        if compressed_jac.shape != (_sparsity.shape[0], ncolors):
             raise ValueError(
                 f"Got an invalid compressed Jacobian shape, which can occur if "
                 f"`fn(x)` is not rank-1 with size matching the number of rows in "
                 f"`sparsity`. Compressed Jacobian shape is {compressed_jac.shape} "
-                f"when `sparsity` has shape {sparsity.shape}."
+                f"when `sparsity` has shape {_sparsity.shape}."
             )
 
         if has_aux:
-            return _expand_jacfwd_jac(compressed_jac, input_coloring, sparsity), aux
+            return _expand_jacfwd_jac(compressed_jac, input_coloring, _sparsity), aux
         else:
-            return _expand_jacfwd_jac(compressed_jac, input_coloring, sparsity)
+            return _expand_jacfwd_jac(compressed_jac, input_coloring, _sparsity)
 
     return jacfwd_fn
 
diff --git a/tests/test_sparsejac.py b/tests/test_sparsejac.py
index 30e4104..67e7e6f 100644
--- a/tests/test_sparsejac.py
+++ b/tests/test_sparsejac.py
@@ -429,5 +429,159 @@ def test_input_connectivity_matches_expected(self):
         onp.testing.assert_array_equal(expected, actual.todense())
 
 
+class JacrevColoringTest(unittest.TestCase):
+    def _diagonal_sparsity(self) -> jsparse.BCOO:
+        return jsparse.BCOO.fromdense(jnp.eye(_SIZE))
+
+    def test_jacrev_coloring_returns_correct_type(self):
+        sparsity = self._diagonal_sparsity()
+        coloring = sparsejac.jacrev_coloring(sparsity)
+        self.assertIsInstance(coloring, sparsejac.ReverseModeColoring)
+
+    def test_jacrev_coloring_fields(self):
+        sparsity = self._diagonal_sparsity()
+        coloring = sparsejac.jacrev_coloring(sparsity)
+        self.assertEqual(coloring.output_coloring.shape, (sparsity.shape[0],))
+        self.assertGreater(coloring.ncolors, 0)
+        self.assertLessEqual(coloring.ncolors, sparsity.shape[0])
+        onp.testing.assert_array_equal(
+            jnp.unique(coloring.output_coloring).size, coloring.ncolors
+        )
+
+    def test_jacrev_with_precomputed_coloring_matches_sparsity_api(self):
+        fn = lambda x: x**2
+        sparsity = self._diagonal_sparsity()
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        expected = sparsejac.jacrev(fn, sparsity)(x)
+        coloring = sparsejac.jacrev_coloring(sparsity)
+        actual = sparsejac.jacrev(fn, coloring=coloring)(x)
+        onp.testing.assert_array_equal(expected.todense(), actual.todense())
+
+    def test_jacrev_with_precomputed_coloring_jit(self):
+        fn = lambda x: x**2
+        sparsity = self._diagonal_sparsity()
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        coloring = sparsejac.jacrev_coloring(sparsity)
+
+        @jax.jit
+        def jacrev_sparse(x, coloring):
+            jacfn = sparsejac.jacrev(fn, coloring=coloring)
+            return jacfn(x)
+        
+        jac = jacrev_sparse(x, coloring)
+        onp.testing.assert_array_equal(jax.jacrev(fn)(x), jac.todense())
+
+    def test_jacrev_neither_sparsity_nor_coloring_raises(self):
+        with self.assertRaisesRegex(ValueError, "Exactly one of"):
+            sparsejac.jacrev(lambda x: x)
+
+    def test_jacrev_both_sparsity_and_coloring_raises(self):
+        sparsity = self._diagonal_sparsity()
+        coloring = sparsejac.jacrev_coloring(sparsity)
+        with self.assertRaisesRegex(ValueError, "not both"):
+            sparsejac.jacrev(lambda x: x, sparsity, coloring=coloring)
+
+    def test_jacrev_coloring_sparsity_shape_validation(self):
+        with self.assertRaisesRegex(ValueError, "`sparsity` must be rank-2"):
+            invalid_sparsity = jsparse.BCOO.fromdense(jnp.ones((5, 5, 5)))
+            sparsejac.jacrev_coloring(invalid_sparsity)
+
+    def test_jacrev_coloring_accepts_scipy_sparse(self):
+        import scipy.sparse as ssparse
+
+        sparsity_scipy = ssparse.eye(_SIZE, format="coo")
+        coloring = sparsejac.jacrev_coloring(sparsity_scipy)
+        self.assertIsInstance(coloring, sparsejac.ReverseModeColoring)
+        self.assertEqual(coloring.output_coloring.shape, (_SIZE,))
+
+    def test_jacrev_with_scipy_coloring_matches_jax_coloring(self):
+        import scipy.sparse as ssparse
+
+        fn = lambda x: x**2
+        sparsity_bcoo = self._diagonal_sparsity()
+        sparsity_scipy = ssparse.eye(_SIZE, format="coo")
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        result_bcoo = sparsejac.jacrev(fn, coloring=sparsejac.jacrev_coloring(sparsity_bcoo))(x)
+        result_scipy = sparsejac.jacrev(fn, coloring=sparsejac.jacrev_coloring(sparsity_scipy))(x)
+        onp.testing.assert_array_equal(result_bcoo.todense(), result_scipy.todense())
+
+
+class JacfwdColoringTest(unittest.TestCase):
+    def _diagonal_sparsity(self) -> jsparse.BCOO:
+        return jsparse.BCOO.fromdense(jnp.eye(_SIZE))
+
+    def test_jacfwd_coloring_returns_correct_type(self):
+        sparsity = self._diagonal_sparsity()
+        coloring = sparsejac.jacfwd_coloring(sparsity)
+        self.assertIsInstance(coloring, sparsejac.ForwardModeColoring)
+
+    def test_jacfwd_coloring_fields(self):
+        sparsity = self._diagonal_sparsity()
+        coloring = sparsejac.jacfwd_coloring(sparsity)
+        self.assertEqual(coloring.input_coloring.shape, (sparsity.shape[1],))
+        self.assertGreater(coloring.ncolors, 0)
+        self.assertLessEqual(coloring.ncolors, sparsity.shape[1])
+        onp.testing.assert_array_equal(
+            jnp.unique(coloring.input_coloring).size, coloring.ncolors
+        )
+
+    def test_jacfwd_with_precomputed_coloring_matches_sparsity_api(self):
+        fn = lambda x: x**2
+        sparsity = self._diagonal_sparsity()
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        expected = sparsejac.jacfwd(fn, sparsity)(x)
+        coloring = sparsejac.jacfwd_coloring(sparsity)
+        actual = sparsejac.jacfwd(fn, coloring=coloring)(x)
+        onp.testing.assert_array_equal(expected.todense(), actual.todense())
+
+    def test_jacfwd_with_precomputed_coloring_jit(self):
+        fn = lambda x: x**2
+        sparsity = self._diagonal_sparsity()
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        coloring = sparsejac.jacfwd_coloring(sparsity)
+
+        @jax.jit
+        def jacfwd_sparse(x, coloring):
+            jacfn = sparsejac.jacfwd(fn, coloring=coloring)
+            return jacfn(x)
+        
+        jac = jacfwd_sparse(x, coloring)
+        onp.testing.assert_array_equal(jax.jacrev(fn)(x), jac.todense())
+
+    def test_jacfwd_neither_sparsity_nor_coloring_raises(self):
+        with self.assertRaisesRegex(ValueError, "Exactly one of"):
+            sparsejac.jacfwd(lambda x: x)
+
+    def test_jacfwd_both_sparsity_and_coloring_raises(self):
+        sparsity = self._diagonal_sparsity()
+        coloring = sparsejac.jacfwd_coloring(sparsity)
+        with self.assertRaisesRegex(ValueError, "not both"):
+            sparsejac.jacfwd(lambda x: x, sparsity, coloring=coloring)
+
+    def test_jacfwd_coloring_sparsity_shape_validation(self):
+        with self.assertRaisesRegex(ValueError, "`sparsity` must be rank-2"):
+            invalid_sparsity = jsparse.BCOO.fromdense(jnp.ones((5, 5, 5)))
+            sparsejac.jacfwd_coloring(invalid_sparsity)
+
+    def test_jacfwd_coloring_accepts_scipy_sparse(self):
+        import scipy.sparse as ssparse
+
+        sparsity_scipy = ssparse.eye(_SIZE, format="coo")
+        coloring = sparsejac.jacfwd_coloring(sparsity_scipy)
+        self.assertIsInstance(coloring, sparsejac.ForwardModeColoring)
+        self.assertEqual(coloring.input_coloring.shape, (_SIZE,))
+
+    def test_jacfwd_with_scipy_coloring_matches_jax_coloring(self):
+        import scipy.sparse as ssparse
+
+        fn = lambda x: x**2
+        sparsity_bcoo = self._diagonal_sparsity()
+        sparsity_scipy = ssparse.eye(_SIZE, format="coo")
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        result_bcoo = sparsejac.jacfwd(fn, coloring=sparsejac.jacfwd_coloring(sparsity_bcoo))(x)
+        result_scipy = sparsejac.jacfwd(fn, coloring=sparsejac.jacfwd_coloring(sparsity_scipy))(x)
+        onp.testing.assert_array_equal(result_bcoo.todense(), result_scipy.todense())
+
+
 if __name__ == "__main__":
     unittest.main(argv=[""], verbosity=2, exit=False)