# 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
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:
        raise RuntimeError(
            "Reverse-mode autodifferentiation is disabled for this operation."
        )


_nondifferentiable_output_p.def_impl(lambda x: x)
_nondifferentiable_output_p.def_abstract_eval(lambda x: x)
batching.primitive_batchers[
    _nondifferentiable_output_p
] = _nondifferentiable_output_batch
xla.register_translation(
    _nondifferentiable_output_p,
    xla.lower_fun(lambda x: x, multiple_results=False, new_style=True),
)
ad.primitive_jvps[_nondifferentiable_output_p] = _nondifferentiable_output_jvp
ad.primitive_transposes[
    _nondifferentiable_output_p
] = _nondifferentiable_output_transpose



# 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

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:
        raise RuntimeError(
            "Reverse-mode autodifferentiation is disabled for this operation."
        )


_nondifferentiable_output_p.def_impl(lambda x: x)
_nondifferentiable_output_p.def_abstract_eval(lambda x: x)
batching.primitive_batchers[
    _nondifferentiable_output_p
] = _nondifferentiable_output_batch
xla.register_translation(
    _nondifferentiable_output_p,
    xla.lower_fun(lambda x: x, multiple_results=False, new_style=True),
)
ad.primitive_jvps[_nondifferentiable_output_p] = _nondifferentiable_output_jvp
ad.primitive_transposes[
    _nondifferentiable_output_p
] = _nondifferentiable_output_transpose



def nondifferentiable_output(x: PyTree) -> PyTree: