# Path: diffrax/custom_types.py
# class PyTree:
#     def __class_getitem__(cls, item):
#         class X:
#             pass
# 
#         X.__module__ = "builtins"
#         X.__qualname__ = _maybe_tuple_to_str(item)
#         return _PyTree[X]


from typing import Any
from ..custom_types import PyTree
import equinox as eqx
import jax
import jax.interpreters.ad as ad
import jax.interpreters.batching as batching
import jax.interpreters.xla as xla
import jax.lax as lax
        return True
    elif isinstance(x, jax.core.Tracer):
        return any(is_perturbed(attr) for name, attr in x._contents())
    else:
        return False


def nondifferentiable_input(x: PyTree, name: str) -> None:
    if any(is_perturbed(xi) for xi in jax.tree_leaves(x)):
        raise ValueError(f"Cannot differentiate {name}.")


_nondifferentiable_output_p = jax.core.Primitive("nondifferentiable_output")


def _nondifferentiable_output_batch(x, batch_axes):
    (x,) = x
    (batch_axes,) = batch_axes
    return nondifferentiable_output(x), batch_axes


def _nondifferentiable_output_jvp(primals, tangents):
    (primals,) = primals
    (tangents,) = tangents
    return nondifferentiable_output(primals), nondifferentiable_output(tangents)


def _nondifferentiable_output_transpose(cts_in, _):
    if isinstance(cts_in, ad.Zero):
        return ad.Zero  # the class, not an instance

# Path: diffrax/custom_types.py
# class PyTree:
#     def __class_getitem__(cls, item):
#         class X:
#             pass
# 
#         X.__module__ = "builtins"
#         X.__qualname__ = _maybe_tuple_to_str(item)
#         return _PyTree[X]


# Path: diffrax/misc/ad.py
from typing import Any
from ..custom_types import PyTree
import equinox as eqx
import jax
import jax.interpreters.ad as ad
import jax.interpreters.batching as batching
import jax.interpreters.xla as xla
import jax.lax as lax

        return True
    elif isinstance(x, jax.core.Tracer):
        return any(is_perturbed(attr) for name, attr in x._contents())
    else:
        return False


def nondifferentiable_input(x: PyTree, name: str) -> None:
    if any(is_perturbed(xi) for xi in jax.tree_leaves(x)):
        raise ValueError(f"Cannot differentiate {name}.")


_nondifferentiable_output_p = jax.core.Primitive("nondifferentiable_output")


def _nondifferentiable_output_batch(x, batch_axes):
    (x,) = x
    (batch_axes,) = batch_axes
    return nondifferentiable_output(x), batch_axes


def _nondifferentiable_output_jvp(primals, tangents):
    (primals,) = primals
    (tangents,) = tangents
    return nondifferentiable_output(primals), nondifferentiable_output(tangents)


def _nondifferentiable_output_transpose(cts_in, _):
    if isinstance(cts_in, ad.Zero):
        return ad.Zero  # the class, not an instance

    else: