# Copyright 2025 Pasteur Labs. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

import re
import types
from collections.abc import Callable, Iterable, Mapping, Sequence
from copy import copy
from typing import (
    Annotated,
    Any,
    ClassVar,
    Literal,
    TypeVar,
    Union,
    get_args,
    get_origin,
)

from pydantic import (
    AfterValidator,
    BaseModel,
    ConfigDict,
    Field,
    RootModel,
    TypeAdapter,
    ValidationError,
    ValidationInfo,
    create_model,
    field_validator,
)

from .schema_types import (
    Array,
    PydanticArrayAnnotation,
    ShapeDType,
    is_differentiable,
    safe_issubclass,
)
from .tree_transforms import get_at_path

# Constants to mark sequence and dict indexing in pytree paths
SEQ_INDEX_SENTINEL = object()
DICT_INDEX_SENTINEL = object()

T = TypeVar("T")

# Python has funnily enough two union types now. See https://github.com/python/cpython/issues/105499
# We check against both for compatibility with older versions of Python.
UNION_TYPES = [Union, types.UnionType]


def _construct_annotated(obj: Any, metadata: Iterable[Any]) -> Any:
    """Construct an Annotated type with the given metadata."""
    out = obj
    # Exploit that repeated Annotated applications are flattened
    for meta in metadata:
        out = Annotated[out, meta]
    return out


def _is_annotated(obj: Any) -> bool:
    """Check if an object is typing.Annotated or typing_extensions.Annotated."""
    return get_origin(obj) is Annotated


def apply_function_to_model_tree(
    Schema: type[BaseModel],
    func: Callable[[type, tuple], type],
    model_prefix: str = "",
    default_model_config: dict[str, Any] | None = None,
) -> type[BaseModel]:
    """Apply a function to all leaves of a Pydantic model, recursing into containers + nested models.

    A leaf is any type annotation that is not a Pydantic model or a container type.

    The given function should take two arguments: the type annotation and the path to it
    in the model tree (as a list of field names). It should return the new type annotation
    for the field, or None to remove the field from the model.

    The notion of path is used to handle nested models and containers. For example, given
    a model like this:

    class MyModel(BaseModel):
        a: list[dict[str, int]]

    The path to the field "a" would be ["a"], and the path to the int type would be
    ["a", SEQ_INDEX_SENTINEL, DICT_INDEX_SENTINEL].
    """
    if default_model_config is None:
        default_model_config = {}

    seen_models = set()

    def _recurse_over_model_tree(treeobj: Any, path: list[str]) -> Any:
        # Get the origin type of the annotation, e.g. List for List[int]
        origin_type = get_origin(treeobj)
        deprecated_types = ["List", "Dict", "Set", "FrozenSet", "Tuple"]

        if origin_type is not None and origin_type.__name__ in deprecated_types:
            # Raise error if the user tries to use a deprecated type
            raise TypeError(
                f"'{origin_type}' annotation deprecated. Use '{origin_type.lower()}' instead."
            )

        if safe_issubclass(treeobj, BaseModel):
            # Recurse into (nested) pydantic models

            if id(treeobj) in seen_models:
                raise ValueError(
                    f"Recursive model definition detected at path {path}: {treeobj.__name__}. "
                    "Recursive models are not supported as Tesseract schemas."
                )
            seen_models.add(id(treeobj))

            new_fields = {}

            for field_name, field in treeobj.model_fields.items():
                field_type = field.annotation
                if field.metadata:
                    field_type = _construct_annotated(field_type, field.metadata)

                field_path = [*path, field_name]
                new_type = _recurse_over_model_tree(field_type, field_path)
                if new_type is None:
                    continue

                new_field = copy(field)
                # Need to strip off metadata to trigger re-evaluation of the field
                new_field.metadata = []

                new_fields[field_name] = (new_type, new_field)

            # We only forbid encountering the same model twice if it is within the same subtree
            seen_models.remove(id(treeobj))

            # Only override model_config if it is not already present
            # in the pydantic model definition
            model_config = ConfigDict(**default_model_config)
            model_config.update(treeobj.model_config)

            # Make sure that reconstructed RootModels don't contain unsupported 'extra' config.
            if safe_issubclass(treeobj, RootModel):
                model_config.pop("extra", None)

            return create_model(
                f"{model_prefix}{treeobj.__name__}",
                **new_fields,
                model_config=(ConfigDict, model_config),
                __base__=treeobj,
            )

        elif _is_annotated(treeobj):
            # Recurse into Annotated types by stripping off right-most (most recent) metadata

            # Using __origin__ and __metadata__ instead of get_args for compatibility with typing_extensions.Annotated
            inner_type = treeobj.__origin__
            *other_meta, current_meta = treeobj.__metadata__

            # Recurse into Annotated[inner_type, *other_meta] and re-attach current_meta to the result
            inner_type = _construct_annotated(inner_type, other_meta)
            new_type = _recurse_over_model_tree(inner_type, path)

            if new_type is None:
                return None

            return Annotated[new_type, current_meta]

        elif any(origin_type is t for t in UNION_TYPES):
            # Recurse into Union and Optional (which is a subclass of Union)
            args = get_args(treeobj)
            newargs = [_recurse_over_model_tree(arg, path) for arg in args]
            newargs = [arg for arg in newargs if arg is not None]
            if not newargs:
                return None

            return Union[tuple(newargs)]  # noqa: UP007

        elif safe_issubclass(origin_type, Mapping):
            # Recurse into dict-likes
            args = get_args(treeobj)
            if not args:
                return origin_type

            key_type, value_type = args

            # Append sentinel to path to indicate dict indexing occurred
            new_path = [*path, DICT_INDEX_SENTINEL]

            new_type = _recurse_over_model_tree(value_type, new_path)
            if new_type is None:
                return None

            return origin_type[key_type, new_type]

        elif safe_issubclass(origin_type, Iterable):
            # Recurse into List, Set, Tuple
            value_type = get_args(treeobj)
            if not value_type:
                return origin_type

            # Append sentinel to path to indicate sequence indexing occurred
            new_path = [*path, SEQ_INDEX_SENTINEL]

            new_types = [_recurse_over_model_tree(vt, new_path) for vt in value_type]
            new_types = [nt for nt in new_types if nt is not None]
            if not new_types:
                return None

            return origin_type[tuple(new_types)]

        # Reached a leaf -> apply func
        return func(treeobj, tuple(path))

    return _recurse_over_model_tree(Schema, [])


def _serialize_diffable_arrays(
    obj: dict[tuple, Any],
) -> dict[str, dict[str, Any]]:
    """Convert a dict {path_tuple: array_type} to a dict {path_str: {shape, dtype}}.

    path_str is a string representation of the path, with SEQ_INDEX_SENTINEL and DICT_INDEX_SENTINEL
    replaced by indexing syntax (e.g. `foo.[].{}`).
    """
    serialized = {}
    for pathtuple, value in obj.items():
        shape = value.expected_shape
        dtype = value.expected_dtype

        # Ensure shape is JSON serializable
        if shape is Ellipsis:
            shape = None
        else:
            shape = tuple(shape)

        # Replace sentinel values with indexing syntax
        str_parts = []
        for part in pathtuple:
            if part is SEQ_INDEX_SENTINEL:
                str_parts.append("[]")
            elif part is DICT_INDEX_SENTINEL:
                str_parts.append("{}")
            else:
                str_parts.append(part)

        serialized[".".join(str_parts)] = {
            "shape": shape,
            "dtype": dtype,
        }

    return serialized


def create_apply_schema(
    InputSchema: type[BaseModel], OutputSchema: type[BaseModel]
) -> tuple[type[BaseModel], type[BaseModel]]:
    """Create the input / output schemas for the /apply endpoint."""
    # We add metadata to the input and output schemas to indicate which fields are differentiable,
    # what their paths are, and which expected shape / dtype they have.
    # This is used internally and by some official clients, but not advertised as part of the public API,
    # so people should not rely on it.
    diffable_input_paths = get_all_model_path_patterns(
        InputSchema, filter_fn=is_differentiable
    )
    diffable_output_paths = get_all_model_path_patterns(
        OutputSchema, filter_fn=is_differentiable
    )

    InputSchema = apply_function_to_model_tree(
        InputSchema,
        lambda x, _: x,
        model_prefix="Apply_",
        default_model_config=dict(extra="forbid"),
    )
    OutputSchema = apply_function_to_model_tree(
        OutputSchema,
        lambda x, _: x,
        model_prefix="Apply_",
        default_model_config=dict(extra="forbid"),
    )

    class ApplyInputSchema(BaseModel):
        inputs: InputSchema = Field(
            ..., description="The input data to apply the Tesseract to."
        )

        differentiable_arrays: ClassVar[dict[str, dict[str, Any]]] = (
            _serialize_diffable_arrays(diffable_input_paths)
        )
        model_config = ConfigDict(
            extra="forbid",
            json_schema_extra={"differentiable_arrays": differentiable_arrays},
        )

    class ApplyOutputSchema(RootModel):
        root: OutputSchema = Field(
            ..., description="The output data from applying the Tesseract."
        )
        differentiable_arrays: ClassVar[dict[str, dict[str, Any]]] = (
            _serialize_diffable_arrays(diffable_output_paths)
        )
        model_config = ConfigDict(
            json_schema_extra={"differentiable_arrays": differentiable_arrays}
        )

    return ApplyInputSchema, ApplyOutputSchema


def create_abstract_eval_schema(
    InputSchema: type[BaseModel], OutputSchema: type[BaseModel]
) -> tuple[type[BaseModel], type[BaseModel]]:
    """Create new Pydantic models that represent abstract versions of the given schemas.

    The returned schemas will have the same structure as the input schemas, but with array
    fields replaced by ShapeDType objects.
    """

    def replace_array_with_shapedtype(obj: T, _: Any) -> T | type[ShapeDType]:
        if safe_issubclass(obj, PydanticArrayAnnotation):
            return ShapeDType.from_array_type(obj)
        return obj

    GeneratedInputSchema = apply_function_to_model_tree(
        InputSchema,
        replace_array_with_shapedtype,
        model_prefix="AbstractEval_",
        default_model_config=dict(extra="forbid"),
    )

    GeneratedOutputSchema = apply_function_to_model_tree(
        OutputSchema,
        replace_array_with_shapedtype,
        model_prefix="AbstractEval_",
        default_model_config=dict(extra="forbid"),
    )

    class AbstractInputSchema(BaseModel):
        inputs: GeneratedInputSchema = Field(
            ...,
            description=(
                "The abstract input data to evaluate the Tesseract on. Has the same structure as InputSchema, "
                "but with array fields replaced by ShapeDType."
            ),
        )
        model_config = ConfigDict(extra="forbid")

    class AbstractOutputSchema(RootModel):
        root: GeneratedOutputSchema = Field(
            ...,
            description=(
                "Abstract outputs with the same structure as OutputSchema, but with array fields "
                "replaced by ShapeDType."
            ),
        )

    return AbstractInputSchema, AbstractOutputSchema


def get_all_model_path_patterns(
    schema: type[BaseModel], filter_fn: Callable[[type], bool] | None = None
) -> dict[tuple, type]:
    """Return a dictionary mapping path patterns to their types.

    For containers (dict, list), the path will include sentinels and map to the container's inner type.
    For model fields, maps to the field's type.
    For leaves (primitives, arrays), maps to the actual type.

    Args:
        schema: The Pydantic model schema to traverse.
        filter_fn: Optional predicate to filter which paths to include.
                   If None, includes all paths. If provided, only includes
                   paths where filter_fn(type) returns True.
    """
    path_to_type = {}

    def add_path(obj: T, path: tuple) -> T:
        if filter_fn is None or filter_fn(obj):
            path_to_type[path] = obj
        return obj

    apply_function_to_model_tree(schema, add_path)
    return path_to_type


def _path_to_pattern(path: Sequence[str | object]) -> str:
    """Return a type describing valid paths for all passed paths."""
    # Check if path includes sequence or dict indexing -- in this case, we can't use the
    # path as a literal and need to use a regex pattern instead.
    final_path = []
    is_literal = True
    for part in path:
        if part is SEQ_INDEX_SENTINEL:
            is_literal = False
            part = r"\[-?\d+\]"
        elif part is DICT_INDEX_SENTINEL:
            is_literal = False
            part = r"\{[\w \-]+\}"
        final_path.append(part)

    if is_literal:
        return ".".join(final_path)

    return r"\.".join(final_path)


def _pattern_to_type(pattern: str) -> type:
    """Convert a string pattern (which may be a literal or regex) to a type that can be used for validation."""
    if _is_regex_pattern(pattern):
        return Annotated[str, Field(pattern=pattern)]
    else:
        return Literal[pattern]


def _is_regex_pattern(pattern: str) -> bool:
    # Poor man's regex detection, but it should work for our purposes
    return "\\" in pattern


def input_path_validator(path: str, info: ValidationInfo) -> str:
    """Validate that the given path points to a valid input key."""
    if "[" in path or "{" in path:
        try:
            get_at_path(info.data["inputs"], path)
        except (LookupError, AttributeError) as exc:
            raise ValueError(
                f"Could not find {info.field_name} path {path} in inputs."
            ) from exc
    return path


def create_gradient_schema(
    InputSchema: type[BaseModel],
    OutputSchema: type[BaseModel],
    gradient_type: Literal["jacobian", "jvp", "vjp"],
) -> tuple[type[BaseModel], type[BaseModel]]:
    """Generate the input / outputs schemas for gradient endpoints like /jacobian, /jacobian_vector_product, etc.

    Returns a tuple (InputSchema, OutputSchema) with the generated schemas.
    """
    diffable_input_paths = get_all_model_path_patterns(
        InputSchema, filter_fn=is_differentiable
    )
    diffable_output_paths = get_all_model_path_patterns(
        OutputSchema, filter_fn=is_differentiable
    )

    if not diffable_input_paths:
        raise RuntimeError("No differentiable inputs found in the input schema")

    if not diffable_output_paths:
        raise RuntimeError("No differentiable outputs found in the output schema")

    diffable_input_patterns = {
        _path_to_pattern(path): obj for path, obj in diffable_input_paths.items()
    }
    diffable_output_patterns = {
        _path_to_pattern(path): obj for path, obj in diffable_output_paths.items()
    }

    diffable_input_type = Union[  # noqa: UP007
        tuple(_pattern_to_type(p) for p in diffable_input_patterns)
    ]
    diffable_output_type = Union[  # noqa: UP007
        tuple(_pattern_to_type(p) for p in diffable_output_patterns)
    ]

    def _find_annotation_from_path(path_patterns: dict, concrete_path: str) -> type:
        """Find the PydanticArrayAnnotation for a given concrete path."""
        for path_pattern, array_type in path_patterns.items():
            if _is_regex_pattern(path_pattern):
                path_matches = bool(re.match(path_pattern, concrete_path))
            else:
                path_matches = path_pattern == concrete_path

            if path_matches:
                return array_type

        raise ValueError(f"Invalid path: {concrete_path}")

    def _find_shape_from_path(path_patterns: dict, concrete_path: str) -> tuple:
        return _find_annotation_from_path(path_patterns, concrete_path).expected_shape

    InputSchema = apply_function_to_model_tree(
        InputSchema,
        lambda x, _: x,
        model_prefix=f"{gradient_type.title()}_",
        default_model_config=dict(extra="forbid"),
    )

    def result_validator(
        result: dict[str, Any], info: ValidationInfo
    ) -> dict[str, Any]:
        """Validate the result of a gradient endpoint computation.

        Since the structure of the result is already validated in core.py, we only need to check the shapes
        to ensure they match what's expected from the schema.
        """
        if gradient_type == "jacobian":
            if set(info.context["output_keys"]) != set(result.keys()):
                raise ValueError(
                    "Error when validating output of jacobian:\n"
                    f"Expected keys {info.context['output_keys']} in output; got {set(result.keys())}"
                )
            for output_path, subout in result.items():
                if set(info.context["input_keys"]) != set(subout.keys()):
                    raise ValueError(
                        "Error when validating output of jacobian:\n"
                        "Expected output with structure "
                        f"{{{tuple(info.context['output_keys'])}: {{{tuple(info.context['input_keys'])}: ...}}}}, "
                        f"got {set(subout.keys())} for output key {output_path}."
                    )
                output_shape = _find_shape_from_path(
                    diffable_output_patterns, output_path
                )
                for input_path, arr in subout.items():
                    input_shape = _find_shape_from_path(
                        diffable_input_patterns, input_path
                    )

                    if output_shape is Ellipsis and input_shape is Ellipsis:
                        # Everything goes
                        continue
                    elif output_shape is Ellipsis:
                        expected_shape = (
                            *input_shape,
                            arr.shape[-len(input_shape) :],
                        )
                    elif input_shape is Ellipsis:
                        expected_shape = (
                            *output_shape,
                            *arr.shape[: len(output_shape)],
                        )
                    else:
                        expected_shape = (*output_shape, *input_shape)

                    expected_annotation = Array[expected_shape, str(arr.dtype)]
                    try:
                        result[output_path][input_path] = TypeAdapter(
                            expected_annotation
                        ).validate_python(arr, context=info.context)
                    except ValidationError as e:
                        raise ValueError(
                            f"Jacobian result [{output_path}][{input_path}]: {e}"
                        ) from e

        elif gradient_type == "jvp":
            if set(info.context["output_keys"]) != set(result.keys()):
                raise ValueError(
                    f"Expected keys {info.context['output_keys']} in output; got {set(result.keys())}"
                )
            for output_path, arr in result.items():
                expected_annotation = _find_annotation_from_path(
                    diffable_output_patterns, output_path
                )
                try:
                    result[output_path] = TypeAdapter(
                        expected_annotation
                    ).validate_python(arr, context=info.context)
                except ValidationError as e:
                    raise ValueError(f"JVP result [{output_path}]: {e}") from e

        elif gradient_type == "vjp":
            if set(info.context["input_keys"]) != set(result.keys()):
                raise ValueError(
                    f"Expected keys {info.context['input_keys']} in output; got {set(result.keys())}"
                )
            for input_path, arr in result.items():
                expected_annotation = _find_annotation_from_path(
                    diffable_input_patterns, input_path
                )
                try:
                    result[input_path] = TypeAdapter(
                        expected_annotation
                    ).validate_python(arr, context=info.context)
                except ValidationError as e:
                    raise ValueError(f"VJP result [{input_path}]: {e}") from e

        return result

    # NOTE (dh): Literal[somevar] is marked as an error by Pylance, but this doesn't throw a runtime error, so I'd
    # prefer to keep this instead of introducing more arcane hacks to create Literal types at runtime.

    if gradient_type == "jacobian":

        class JacobianInputSchema(BaseModel):
            inputs: InputSchema = Field(
                ..., description="The input data to compute the Jacobian at."
            )
            jac_inputs: set[
                Annotated[diffable_input_type, AfterValidator(input_path_validator)]
            ] = Field(
                ...,
                description="The set of differentiable inputs to compute the Jacobian with respect to.",
                min_length=1,
            )
            jac_outputs: set[diffable_output_type] = Field(
                ...,
                description="The set of differentiable outputs to compute the Jacobian of.",
                min_length=1,
            )

            model_config = ConfigDict(extra="forbid")

        class JacobianOutputSchema(RootModel):
            root: Annotated[
                dict[diffable_output_type, dict[diffable_input_type, Array[..., None]]],
                AfterValidator(result_validator),
            ] = Field(
                ...,
                description=(
                    "Container for the results of Jacobian computations. "
                    "The result represents a nested structure of the Jacobian matrix "
                    "as a mapping with structure ``{jac_outputs: {jac_inputs: array}}``. "
                    "The shape of each array is the concatenation of the shapes of the output and input arrays, "
                    "i.e. ``(*output_array.shape, *input_array.shape)``."
                ),
            )

        InputSchema = JacobianInputSchema
        OutputSchema = JacobianOutputSchema

    elif gradient_type == "jvp":

        class JVPInputSchema(BaseModel):
            inputs: InputSchema = Field(
                ..., description="The input data to compute the JVP at."
            )
            jvp_inputs: set[
                Annotated[diffable_input_type, AfterValidator(input_path_validator)]
            ] = Field(
                ...,
                description="The set of differentiable inputs to compute the JVP with respect to.",
                min_length=1,
            )
            jvp_outputs: set[diffable_output_type] = Field(
                ...,
                description="The set of differentiable outputs to compute the JVP of.",
                min_length=1,
            )
            tangent_vector: dict[diffable_input_type, Array[..., None]] = Field(
                ...,
                description=(
                    "Tangent vector to multiply the Jacobian with. "
                    "Expected to be a mapping with structure ``{jvp_inputs: array}``. "
                    "The shape of each array is the same as the shape of the corresponding input array."
                ),
            )
            model_config = ConfigDict(extra="forbid")

            @field_validator("tangent_vector", mode="after")
            @classmethod
            def validate_tangent_vector(
                cls, tangent_vector: dict, info: ValidationInfo
            ) -> dict:
                """Validate tangent vector keys, shapes, and dtypes."""
                try:
                    if set(tangent_vector.keys()) != info.data["jvp_inputs"]:
                        raise ValueError(
                            f"Expected tangent vector with keys conforming to jvp_inputs: {info.data['jvp_inputs']}, "
                            f"got {set(tangent_vector.keys())}."
                        )
                except KeyError as e:
                    raise ValueError(
                        "Unable to validate tangent vector as jvp_inputs either missing or invalid"
                    ) from e

                validated = {}
                for path, arr in tangent_vector.items():
                    ref_arr = get_at_path(info.data["inputs"], path)
                    if hasattr(ref_arr, "shape"):
                        ref_shape = ref_arr.shape
                    elif isinstance(ref_arr, dict):
                        ref_shape = tuple(ref_arr["shape"])
                    else:
                        # Scalar (e.g. Differentiable[Float32])
                        ref_shape = ()
                    annotation = _find_annotation_from_path(
                        diffable_input_patterns, path
                    )
                    exact_annotation = Array[ref_shape, annotation.expected_dtype]
                    try:
                        validated[path] = TypeAdapter(exact_annotation).validate_python(
                            arr, context=info.context
                        )
                    except ValidationError as e:
                        raise ValueError(f"Tangent vector '{path}': {e}") from e
                return validated

        class JVPOutputSchema(RootModel):
            root: Annotated[
                dict[diffable_output_type, Array[..., None]],
                AfterValidator(result_validator),
            ] = Field(
                ...,
                description=(
                    "Container for the results of Jacobian-vector products. "
                    "The result is a mapping with structure ``{jvp_outputs: array}``. "
                    "The shape of each array is the same as the shape of the corresponding output array."
                ),
            )

        InputSchema = JVPInputSchema
        OutputSchema = JVPOutputSchema

    elif gradient_type == "vjp":

        class VJPInputSchema(BaseModel):
            inputs: InputSchema = Field(
                ..., description="The input data to compute the VJP at."
            )
            vjp_inputs: set[
                Annotated[diffable_input_type, AfterValidator(input_path_validator)]
            ] = Field(
                ...,
                description="The set of differentiable inputs to compute the VJP with respect to.",
                min_length=1,
            )
            vjp_outputs: set[diffable_output_type] = Field(
                ...,
                description="The set of differentiable outputs to compute the VJP of.",
                min_length=1,
            )
            cotangent_vector: dict[diffable_output_type, Array[..., None]] = Field(
                ...,
                description=(
                    "Cotangent vector to multiply the Jacobian with. "
                    "Expected to be a mapping with structure ``{vjp_outputs: array}``. "
                    "The shape of each array is the same as the shape of the corresponding output array."
                ),
            )
            model_config = ConfigDict(extra="forbid")

            @field_validator("cotangent_vector", mode="after")
            @classmethod
            def validate_cotangent_vector(
                cls, cotangent_vector: dict, info: ValidationInfo
            ) -> dict:
                """Validate cotangent vector keys, shapes, and dtypes."""
                try:
                    if set(cotangent_vector.keys()) != info.data["vjp_outputs"]:
                        raise ValueError(
                            "Expected cotangent vector with keys conforming to vjp_outputs: "
                            f"{info.data['vjp_outputs']}, got {set(cotangent_vector.keys())}."
                        )
                except KeyError as e:
                    raise ValueError(
                        "Unable to validate cotangent vector as vjp_outputs either missing or invalid"
                    ) from e

                validated = {}
                for path, arr in cotangent_vector.items():
                    annotation = _find_annotation_from_path(
                        diffable_output_patterns, path
                    )
                    try:
                        validated[path] = TypeAdapter(annotation).validate_python(
                            arr, context=info.context
                        )
                    except ValidationError as e:
                        raise ValueError(f"Cotangent vector '{path}': {e}") from e
                return validated

        class VJPOutputSchema(RootModel):
            root: Annotated[
                dict[diffable_input_type, Array[..., None]],
                AfterValidator(result_validator),
            ] = Field(
                ...,
                description=(
                    "Container for the results of vector-Jacobian products. "
                    "The result is a mapping with structure ``{vjp_inputs: array}``. "
                    "The shape of each array is the same as the shape of the corresponding input array."
                ),
            )

        InputSchema = VJPInputSchema
        OutputSchema = VJPOutputSchema

    return InputSchema, OutputSchema
