Repository: pasteurlabs/tesseract-core

Issue:
Make input validation message less scary
Pydantic validation spits out a really scary message that is not easy to parse. To some extent this is necessary of course, but maybe we could at least hide parts of it, e.g. stack trace that carries no useful info?

To illustrate, let's use `tesseract init --name test-init --recipe jax` as an example. Here is the input that works:

```
❯ cat example_input.json                                                                  
{
    "inputs":{
       "example":{
          "object_type":"array",
          "shape":[],
          "dtype":"float32",
          "data":{
             "buffer":[1],
             "encoding":"json"
          }
       }
    }
}
❯ tesseract run test-init apply @example_input.json                           
{"example":{"object_type":"array","shape":[],"dtype":"float32","data":{"buffer":1.0,"encoding":"json"}}}
```

But if we change `"shape":[],` to `"shape":[1],`, it gets very upset:

```
❯ tesseract run test-init apply @example_input.json                                 
 [-] Error running Tesseract. 

Command 'apply @/mnt/payload.json' in image 'test-init' returned non-zero exit status 1: 
Traceback (most recent call last):
  File "/python-env/bin/tesseract-runtime", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/python-env/lib/python3.12/site-packages/tesseract_core/runtime/cli.py", line 323, in main
    cli(auto_envvar_prefix="TESSERACT_RUNTIME")
  File "/python-env/lib/python3.12/site-packages/click/core.py", line 1161, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python-env/lib/python3.12/site-packages/typer/core.py", line 743, in main
    return _main(
           ^^^^^^
  File "/python-env/lib/python3.12/site-packages/typer/core.py", line 198, in _main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/python-env/lib/python3.12/site-packages/click/core.py", line 1697, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python-env/lib/python3.12/site-packages/click/core.py", line 1443, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python-env/lib/python3.12/site-packages/click/core.py", line 788, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python-env/lib/python3.12/site-packages/tesseract_core/runtime/cli.py", line 226, in _callback_wrapper
    user_function_args["payload"] = InputSchema.model_validate(
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python-env/lib/python3.12/site-packages/pydantic/main.py", line 627, in model_validate
    return cls.__pydantic_validator__.validate_python(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.ValidationError: 2 validation errors for ApplyInputSchema
inputs.example.`chain[EncodedArrayModel__scalar__float32__DIFFERENTIABLE,function-plain[functools.partial(<function decode_array at 0xffff7fd19440>, expected_shape=(), expected_dtype='float32')()]]`.shape
  Tuple should have at most 0 items after validation, not 1 [type=too_long, input_value=[1], input_type=list]
    For further information visit https://errors.pydantic.dev/2.10/v/too_long
inputs.example.`function-plain[functools.partial(<function python_to_array at 0xffff7fd19300>, expected_shape=(), expected_dtype='float32')()]`
  Value error, Could not convert {'object_type': 'array', 'shape': [1], 'dtype': 'float32', 'data': {'buffer': [1], 'encoding': 'json'}} to NumPy array [type=value_error, input_value={'object_type': 'array', ...1], 'encoding': 'json'}}, input_type=dict]
    For further information visit https://errors.pydantic.dev/2.10/v/value_error

 [x] Aborting
```



Discussion:
jpbrodrick89: Dependably bump seems to have broken some file hashes not sure how to fix.
jpbrodrick89: Re-ran the couple of failing tests and they passed this time. Maybe itwas something to do with pip version (the passed test don't have a warning to update pip). 🤔 

Patch:
diff --git a/examples/dataloader/tesseract_api.py b/examples/dataloader/tesseract_api.py
index c80eaf6..0356b58 100644
--- a/examples/dataloader/tesseract_api.py
+++ b/examples/dataloader/tesseract_api.py
@@ -57,7 +57,7 @@ def jacobian(
         assert len(key_parts) == 2
 
         idx = int(key_parts[1][1:-1])
-        jac = np.ones(3, *inputs.data[idx].shape)
+        jac = np.zeros((3, *inputs.data[idx].shape))
         jac[0, :, 0] = 1.0
         jac[1, :, 1] = 1.0
         jac[2, :, 2] = 1.0
diff --git a/examples/vectoradd_jax/example_bad_inputs.json b/examples/vectoradd_jax/example_bad_inputs.json
new file mode 100644
index 0000000..158b7e1
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_inputs.json
@@ -0,0 +1,19 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+            ],
+           "x": 3.0
+       }
+    }
+}
diff --git a/examples/vectoradd_jax/example_bad_inputs2.json b/examples/vectoradd_jax/example_bad_inputs2.json
new file mode 100644
index 0000000..03ac030
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_inputs2.json
@@ -0,0 +1,27 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+             ],
+            "s":{
+              "object_type":"array",
+              "shape":[1],
+              "dtype":"float32",
+              "data":{
+                "buffer":[1],
+                "encoding":"json"
+              }
+            }
+        }
+    }
+}
diff --git a/examples/vectoradd_jax/example_bad_vjp_inputs.json b/examples/vectoradd_jax/example_bad_vjp_inputs.json
new file mode 100644
index 0000000..73e44c0
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_vjp_inputs.json
@@ -0,0 +1,38 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+            ],
+           "s": 3.0
+       }
+    },
+    "vjp_inputs": [
+        "a.v",
+        "b.v"
+    ],
+    "vjp_outputs": [
+        "vector_add.result"
+    ],
+    "cotangent_vector": {
+        "vector_add.result": [
+            0.1,
+            0.2,
+            0.3
+        ],
+        "vector_add.normed_result": [
+            0.4,
+            0.5,
+            0.6
+        ]
+    }
+}
diff --git a/examples/vectoradd_jax/example_bad_vjp_inputs2.json b/examples/vectoradd_jax/example_bad_vjp_inputs2.json
new file mode 100644
index 0000000..078cd7c
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_vjp_inputs2.json
@@ -0,0 +1,39 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+            ],
+           "s": 3.0
+       }
+    },
+    "vjp_inputs": [
+        "a.v",
+        "b.v"
+    ],
+    "vjp_outputs": [
+        "vector_add.refsult",
+        "vector_add.normed_result"
+    ],
+    "cotangent_vector": {
+        "vector_add.refsult": [
+            0.1,
+            0.2,
+            0.3
+        ],
+        "vector_add.normed_result": [
+            0.4,
+            0.5,
+            0.6
+        ]
+    }
+}
diff --git a/examples/vectoradd_jax/example_bad_vjp_inputs3.json b/examples/vectoradd_jax/example_bad_vjp_inputs3.json
new file mode 100644
index 0000000..c25e711
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_vjp_inputs3.json
@@ -0,0 +1,27 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+            ],
+           "s": 3.0
+       }
+    },
+    "vjp_inputs": [
+        "a.v",
+        "b.v"
+    ],
+    "vjp_outputs": [
+    ],
+    "cotangent_vector": {
+    }
+}
diff --git a/examples/vectoradd_jax/example_bad_vjp_inputs4.json b/examples/vectoradd_jax/example_bad_vjp_inputs4.json
new file mode 100644
index 0000000..361cffa
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_vjp_inputs4.json
@@ -0,0 +1,37 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+            ],
+           "s": 3.0
+       }
+    },
+    "vjp_inputs": [
+    ],
+    "vjp_outputs": [
+        "vector_add.result",
+        "vector_add.normed_result"
+    ],
+    "cotangent_vector": {
+        "vector_add.result": [
+            0.1,
+            0.2,
+            0.3
+        ],
+        "vector_add.normed_result": [
+            0.4,
+            0.5,
+            0.6
+        ]
+    }
+}
diff --git a/examples/vectoradd_jax/example_bad_vjp_inputs5.json b/examples/vectoradd_jax/example_bad_vjp_inputs5.json
new file mode 100644
index 0000000..1d80041
--- /dev/null
+++ b/examples/vectoradd_jax/example_bad_vjp_inputs5.json
@@ -0,0 +1,23 @@
+{
+    "inputs": {
+        "a": {
+            "v": [
+                1.0,
+                2.0,
+                3.0
+            ]
+        },
+        "b": {
+            "v": [
+                4.0,
+                5.0,
+                6.0
+            ],
+           "s": 3.0
+       }
+    },
+    "vjp_inputs": [
+    ],
+    "cotangent_vector": {
+    }
+}
diff --git a/tesseract_core/runtime/cli.py b/tesseract_core/runtime/cli.py
index 0d082a2..9215a9a 100644
--- a/tesseract_core/runtime/cli.py
+++ b/tesseract_core/runtime/cli.py
@@ -14,6 +14,7 @@ from typing import Any, Optional
 
 import click
 import typer
+from pydantic import ValidationError
 
 from tesseract_core.runtime.config import get_config
 from tesseract_core.runtime.core import create_endpoints, get_tesseract_api
@@ -223,9 +224,16 @@ def _create_user_defined_cli_command(
 
         if InputSchema is not None:
             payload, base_dir = optional_args["payload"]
-            user_function_args["payload"] = InputSchema.model_validate(
-                payload, context={"base_dir": base_dir}
-            )
+            try:
+                user_function_args["payload"] = InputSchema.model_validate(
+                    payload, context={"base_dir": base_dir}
+                )
+            except ValidationError as e:
+                raise click.BadParameter(
+                    e,
+                    param=InputSchema,
+                    param_hint=InputSchema.__name__,
+                ) from e
 
         result = user_function(**user_function_args)
         result = output_to_bytes(result, output_format, output_path)
diff --git a/tesseract_core/runtime/core.py b/tesseract_core/runtime/core.py
index 5cc6437..93d600d 100644
--- a/tesseract_core/runtime/core.py
+++ b/tesseract_core/runtime/core.py
@@ -2,13 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import importlib
-import re
 from collections.abc import Callable
 from pathlib import Path
 from types import ModuleType
-from typing import Any, Optional, Union
-
-from pydantic import BaseModel
+from typing import Any, Union
 
 from .config import get_config
 from .schema_generation import (
@@ -16,7 +13,6 @@ from .schema_generation import (
     create_apply_schema,
     create_autodiff_schema,
 )
-from .tree_transforms import get_at_path
 
 
 def load_module_from_path(path: Union[Path, str]) -> ModuleType:
@@ -59,98 +55,6 @@ def get_tesseract_api() -> ModuleType:
     return load_module_from_path(get_config().tesseract_api_path)
 
 
-def _validate_ad_input(
-    inputs: BaseModel,
-    input_keys: set[str],
-    output_keys: set[str],
-    endpoint_name: str,
-    tangent_cotangent_vector: Optional[dict[str, Any]] = None,
-):
-    """Raise an exception if the input of an autodiff function does not conform to given input keys."""
-    # Could be moved to a model validator
-    for input_key in input_keys:
-        if not re.match(r"^[a-zA-Z0-9_.\[\]{}]+$", input_key):
-            raise RuntimeError(
-                f"Error when validating input of function {endpoint_name}:\n"
-                f"Invalid input key {input_key}."
-            )
-
-        try:
-            get_at_path(inputs, input_key)
-        except Exception as exc:
-            raise RuntimeError(
-                f"Error when validating input of function {endpoint_name}:\n"
-                f"Could not find input path {input_key} in inputs."
-            ) from exc
-
-    if endpoint_name == "jacobian_vector_product":
-        # Tangent vector needs same keys as input_keys
-        if set(tangent_cotangent_vector.keys()) != input_keys:
-            raise RuntimeError(
-                f"Error when validating input of function {endpoint_name}:\n"
-                f"Expected tangent vector with keys {input_keys}, got {set(tangent_cotangent_vector.keys())}."
-            )
-
-    if endpoint_name == "vector_jacobian_product":
-        # Cotangent vector needs same keys as output_keys
-        if set(tangent_cotangent_vector.keys()) != output_keys:
-            raise RuntimeError(
-                f"Error when validating input of function {endpoint_name}:\n"
-                f"Expected cotangent vector with keys {output_keys}, got {set(tangent_cotangent_vector.keys())}."
-            )
-
-
-def _validate_ad_output(
-    input: Any,
-    output: Any,
-    input_keys: set[str],
-    output_keys: set[str],
-    endpoint_name: str,
-):
-    """Raise an exception if the output structure of an autodiff function does not conform to given keys."""
-    if not isinstance(output, dict):
-        raise RuntimeError(
-            f"Error when validating output of function {endpoint_name}:\n"
-            f"Expected output to be a dictionary, got {type(output)}"
-        )
-
-    if endpoint_name == "jacobian":
-        if output_keys != output.keys():
-            raise RuntimeError(
-                "Error when validating output of jacobian:\n"
-                f"Expected keys {output_keys} in output; got {set(output.keys())}"
-            )
-        for subkey, subout in output.items():
-            if not isinstance(subout, dict):
-                raise RuntimeError(
-                    "Error when validating output of jacobian:\n"
-                    f"Expected output with structure {{{tuple(output_keys)}: {{{tuple(input_keys)}: ...}}}}, "
-                    f"got unexpected type {type(subout)} for output key {subkey}."
-                )
-            if input_keys != subout.keys():
-                raise RuntimeError(
-                    "Error when validating output of jacobian:\n"
-                    f"Expected output with structure {{{tuple(output_keys)}: {{{tuple(input_keys)}: ...}}}}, "
-                    f"got {set(subout.keys())} for output key {subkey}."
-                )
-
-    elif endpoint_name == "jacobian_vector_product":
-        if output_keys != output.keys():
-            raise RuntimeError(
-                "Error when validating output of jacobian_vector_product:\n"
-                f"Expected keys {output_keys} in output; got {set(output.keys())}"
-            )
-
-    elif endpoint_name == "vector_jacobian_product":
-        if input_keys != output.keys():
-            raise RuntimeError(
-                "Error when validating output of vector_jacobian_product:\n"
-                f"Expected keys {input_keys} in output; got {set(output.keys())}"
-            )
-    else:
-        raise RuntimeError(f"Unknown endpoint name {endpoint_name}")
-
-
 def create_endpoints(api_module: ModuleType) -> list[Callable]:
     """Create the Tesseract API endpoints.
 
@@ -207,16 +111,16 @@ def create_endpoints(api_module: ModuleType) -> list[Callable]:
 
             Differentiates ``jac_outputs`` with respect to ``jac_inputs``, at the point ``inputs``.
             """
-            _validate_ad_input(
-                payload.inputs, payload.jac_inputs, payload.jac_outputs, "jacobian"
-            )
             out = api_module.jacobian(
                 payload.inputs, payload.jac_inputs, payload.jac_outputs
             )
-            _validate_ad_output(
-                payload.inputs, out, payload.jac_inputs, payload.jac_outputs, "jacobian"
+            return JacobianOutputSchema.model_validate(
+                out,
+                context={
+                    "output_keys": payload.jac_outputs,
+                    "input_keys": payload.jac_inputs,
+                },
             )
-            return JacobianOutputSchema.model_validate(out)
 
         endpoints.append(jacobian)
 
@@ -232,27 +136,15 @@ def create_endpoints(api_module: ModuleType) -> list[Callable]:
             Evaluates the Jacobian vector product between the Jacobian given by ``jvp_outputs``
             with respect to ``jvp_inputs`` at the point ``inputs`` and the given tangent vector.
             """
-            _validate_ad_input(
-                payload.inputs,
-                payload.jvp_inputs,
-                payload.jvp_outputs,
-                "jacobian_vector_product",
-                tangent_cotangent_vector=payload.tangent_vector,
-            )
             out = api_module.jacobian_vector_product(
                 payload.inputs,
                 payload.jvp_inputs,
                 payload.jvp_outputs,
                 payload.tangent_vector,
             )
-            _validate_ad_output(
-                payload.inputs,
-                out,
-                payload.jvp_inputs,
-                payload.jvp_outputs,
-                "jacobian_vector_product",
+            return JVPOutputSchema.model_validate(
+                out, context={"output_keys": payload.jvp_outputs}
             )
-            return JVPOutputSchema.model_validate(out)
 
         endpoints.append(jacobian_vector_product)
 
@@ -268,27 +160,15 @@ def create_endpoints(api_module: ModuleType) -> list[Callable]:
             Computes the vector Jacobian product between the Jacobian given by ``vjp_outputs``
             with respect to ``vjp_inputs`` at the point ``inputs`` and the given cotangent vector.
             """
-            _validate_ad_input(
-                payload.inputs,
-                payload.vjp_inputs,
-                payload.vjp_outputs,
-                "vector_jacobian_product",
-                tangent_cotangent_vector=payload.cotangent_vector,
-            )
             out = api_module.vector_jacobian_product(
                 payload.inputs,
                 payload.vjp_inputs,
                 payload.vjp_outputs,
                 payload.cotangent_vector,
             )
-            _validate_ad_output(
-                payload.inputs,
-                out,
-                payload.vjp_inputs,
-                payload.vjp_outputs,
-                "vector_jacobian_product",
+            return VJPOutputSchema.model_validate(
+                out, context={"input_keys": payload.vjp_inputs}
             )
-            return VJPOutputSchema.model_validate(out)
 
         endpoints.append(vector_jacobian_product)
 
diff --git a/tesseract_core/runtime/schema_generation.py b/tesseract_core/runtime/schema_generation.py
index 92c6b9a..cf94f1a 100644
--- a/tesseract_core/runtime/schema_generation.py
+++ b/tesseract_core/runtime/schema_generation.py
@@ -22,7 +22,9 @@ from pydantic import (
     ConfigDict,
     Field,
     RootModel,
+    ValidationInfo,
     create_model,
+    field_validator,
 )
 
 from .schema_types import (
@@ -34,6 +36,7 @@ from .schema_types import (
     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()
@@ -320,10 +323,10 @@ def _path_to_pattern(path: Sequence[Union[str, object]]) -> str:
     for part in path:
         if part is SEQ_INDEX_SENTINEL:
             is_literal = False
-            part = r"\[\d+\]"
+            part = r"\[-?\d+\]"
         elif part is DICT_INDEX_SENTINEL:
             is_literal = False
-            part = r"\{.*?\}"
+            part = r"\{[\w \-]+\}"
         final_path.append(part)
 
     if is_literal:
@@ -336,8 +339,8 @@ 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)]
-
-    return Literal[pattern]
+    else:
+        return Literal[pattern]
 
 
 def _is_regex_pattern(pattern: str) -> bool:
@@ -345,6 +348,17 @@ def _is_regex_pattern(pattern: str) -> bool:
     return "\\" in pattern
 
 
+def input_path_validator(path: str, info: ValidationInfo) -> str:
+    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_autodiff_schema(
     InputSchema: type[BaseModel],
     OutputSchema: type[BaseModel],
@@ -390,7 +404,7 @@ def create_autodiff_schema(
     )
 
     def result_validator(
-        result: dict[str, Any],
+        result: dict[str, Any], info: ValidationInfo
     ) -> dict[str, Any]:
         """Validate the result of a Jacobian computation.
 
@@ -398,7 +412,19 @@ def create_autodiff_schema(
         to ensure they match what's expected from the schema.
         """
         if ad_flavor == "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_outputs, output_path)
                 for input_path, arr in subout.items():
                     input_shape = _find_shape_from_path(diffable_inputs, input_path)
@@ -442,6 +468,10 @@ def create_autodiff_schema(
                         )
 
         elif ad_flavor == "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():
                 output_shape = _find_shape_from_path(diffable_outputs, output_path)
                 if output_shape is Ellipsis:
@@ -465,6 +495,10 @@ def create_autodiff_schema(
                     )
 
         elif ad_flavor == "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():
                 input_shape = _find_shape_from_path(diffable_inputs, input_path)
                 if input_shape is Ellipsis:
@@ -498,7 +532,9 @@ def create_autodiff_schema(
             inputs: InputSchema = Field(
                 ..., description="The input data to compute the Jacobian at."
             )
-            jac_inputs: set[diffable_input_type] = Field(
+            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,
@@ -535,7 +571,9 @@ def create_autodiff_schema(
             inputs: InputSchema = Field(
                 ..., description="The input data to compute the JVP at."
             )
-            jvp_inputs: set[diffable_input_type] = Field(
+            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,
@@ -553,8 +591,28 @@ def create_autodiff_schema(
                     "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:
+                """Raise an exception if the input of an autodiff function does not conform to given input keys."""
+                # Cotangent vector needs same keys as output_keys
+                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
+                return tangent_vector
+
         class JVPOutputSchema(RootModel):
             root: Annotated[
                 dict[diffable_output_type, Array[..., None]],
@@ -577,7 +635,9 @@ def create_autodiff_schema(
             inputs: InputSchema = Field(
                 ..., description="The input data to compute the VJP at."
             )
-            vjp_inputs: set[diffable_input_type] = Field(
+            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,
@@ -598,6 +658,24 @@ def create_autodiff_schema(
 
             model_config = ConfigDict(extra="forbid")
 
+            @field_validator("cotangent_vector", mode="after")
+            @classmethod
+            def validate_cotangent_vector(
+                cls, cotangent_vector: dict, info: ValidationInfo
+            ) -> dict:
+                """Raise an exception if cotangent vector keys are not identical to vjp_outputs."""
+                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
+                return cotangent_vector
+
         class VJPOutputSchema(RootModel):
             root: Annotated[
                 dict[diffable_input_type, Array[..., None]],


Test patch:
diff --git a/tests/endtoend_tests/test_examples.py b/tests/endtoend_tests/test_examples.py
index dae5248..9144364 100644
--- a/tests/endtoend_tests/test_examples.py
+++ b/tests/endtoend_tests/test_examples.py
@@ -11,6 +11,7 @@ Add test cases for specific unit Tesseracts to the TEST_CASES dictionary.
 
 import base64
 import json
+import re
 import time
 import traceback
 from dataclasses import dataclass
@@ -26,6 +27,12 @@ from common import build_tesseract
 from typer.testing import CliRunner
 
 
+# Only necessary when matching multi-word string
+def format_stderr(stderr: str) -> str:
+    no_color = re.sub(r"\x1b\[[0-9;]*m", "", stderr)
+    return " ".join(re.sub(r"[^\w \d_.,!?:;\-]+", " ", no_color).split())
+
+
 def json_normalize(obj: str):
     """Normalize JSON str for comparison."""
     return json.dumps(json.loads(obj), separators=(",", ":"))
@@ -652,7 +659,10 @@ def test_unit_tesseract_endtoend(
             else:
                 assert result.exit_code != 0
                 # Result is an error message
-                output = "\n".join(traceback.format_exception(*result.exc_info))
+                output = "".join(traceback.format_exception(*result.exc_info))
+                # Click with rich adds color codes and boxes to the output
+                # which we need to remove in case they break multi-word pattern matching
+                output = format_stderr(output)
 
             if request.output_contains_pattern is not None:
                 patterns = request.output_contains_pattern
diff --git a/tests/runtime_tests/test_core.py b/tests/runtime_tests/test_core.py
index b360a8b..258e123 100644
--- a/tests/runtime_tests/test_core.py
+++ b/tests/runtime_tests/test_core.py
@@ -11,7 +11,7 @@ from types import ModuleType
 
 import numpy as np
 import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 from tesseract_core.runtime import Array, Differentiable, Float32
 from tesseract_core.runtime.core import create_endpoints
@@ -264,12 +264,12 @@ def test_ad_endpoint(testmodule, ad_inout, endpoint_name):
         ),
         # Non-existent container index
         (
-            "Could not find input path array_seq.[100]",
+            "array_seq.[100]",
             {"array_seq.[100]", "array_dict.{a}", "scalar_diff"},
             {"result_seq.[0]"},
         ),
         (
-            "Could not find input path array_dict.{xyz}",
+            "array_dict.{xyz}",
             {"array_seq.[0]", "array_dict.{xyz}", "scalar_diff"},
             {"result_seq.[0]"},
         ),
@@ -357,7 +357,7 @@ def test_ad_endpoint_invalid(testmodule, ad_inout_invalid, endpoint_name):
             "cotangent_vector": cotangent_vector,
         }
 
-    with pytest.raises((ValueError, RuntimeError)) as excinfo:
+    with pytest.raises(ValidationError) as excinfo:
         inputs = EndpointSchema.model_validate(inputs)
         endpoint_func(inputs)
 
@@ -415,7 +415,7 @@ def test_ad_endpoint_bad_tangent(testmodule, endpoint_name, failure_mode):
             "cotangent_vector": cotangent_vector,
         }
 
-    with pytest.raises((ValueError, RuntimeError)) as excinfo:
+    with pytest.raises(ValidationError) as excinfo:
         inputs = EndpointSchema.model_validate(inputs)
         endpoint_func(inputs)
 
diff --git a/tests/runtime_tests/test_runtime_cli.py b/tests/runtime_tests/test_runtime_cli.py
index 9c5d2cb..c1f9161 100644
--- a/tests/runtime_tests/test_runtime_cli.py
+++ b/tests/runtime_tests/test_runtime_cli.py
@@ -4,6 +4,7 @@
 import copy
 import json
 import os
+import re
 import subprocess
 import sys
 import traceback
@@ -22,6 +23,14 @@ from tesseract_core.runtime.cli import _add_user_commands_to_cli
 from tesseract_core.runtime.cli import tesseract_runtime as cli_cmd
 from tesseract_core.runtime.file_interactions import load_bytes, output_to_bytes
 
+
+# Remove color codes and box
+# Only necessary when matching multi-word string
+def format_stderr(stderr: str) -> str:
+    no_color = re.sub(r"\x1b\[[0-9;]*m", "", stderr)
+    return " ".join(re.sub(r"[^\w \d_.,!?:;\-]+", " ", no_color).split())
+
+
 test_input = {
     "a": [1.0, 2.0, 3.0],
     "b": [1, 1, 1],
@@ -232,12 +241,17 @@ def test_apply_command_binref(cli, cli_runner, dummy_tesseract_module, tmpdir):
     expected = dummy_tesseract_module.apply(test_input_val).model_dump_json()
     assert json.loads(result.stdout) == json.loads(expected)
 
-    with pytest.raises(ValueError, match="binref encoded with a relative path"):
-        result = cli_runner.invoke(
-            cli,
-            ["apply", json.dumps({"inputs": test_input_binref})],
-            catch_exceptions=False,
-        )
+    result = cli_runner.invoke(
+        cli,
+        ["apply", json.dumps({"inputs": test_input_binref})],
+        catch_exceptions=False,
+    )
+    # Click with rich adds color codes and boxes to the output
+    # which we need to remove in case they break multi-word pattern matching
+    stderr_fmt = format_stderr(result.stderr)
+    assert result.exit_code == 2
+    assert "Value error" in stderr_fmt
+    assert "binref encoded with a relative path" in stderr_fmt
 
 
 def test_apply_command_noenv(
@@ -443,7 +457,10 @@ def test_optional_arguments_stay_optional_in_cli(
 
 def test_apply_fails_if_required_args_missing(cli, cli_runner):
     result = cli_runner.invoke(cli, ["apply", json.dumps({"inputs": {"a": [1, 2, 3]}})])
-    assert result.exit_code == 1, result.stdout
+    print("Result:", result)
+    print(result.exit_code, result.stdout, result.stderr)
+    assert result.exit_code == 2
+    assert "missing" in result.stderr
 
 
 def test_stdout_redirect_cli():
diff --git a/tests/runtime_tests/test_schema_generation.py b/tests/runtime_tests/test_schema_generation.py
index 6c49baa..0026714 100644
--- a/tests/runtime_tests/test_schema_generation.py
+++ b/tests/runtime_tests/test_schema_generation.py
@@ -53,7 +53,7 @@ def replace_arrays(inpobj, replace_fn):
 testinput = NestedModel(
     testdiffarr=make_array((5, 6), "float64"),
     testfoo=[SubModel(foo=1.0, bar=[1, 2, 3]), SubModel(foo=4.0, bar=[])],
-    testbar={"hey there!": SubModel(foo=5.0, bar=[2])},
+    testbar={"hey there": SubModel(foo=5.0, bar=[2])},
     testbaz=make_array((1, 2, 3), "uint8"),
     testset={1, 2, 3},
     testtuple=(1, "hello"),
@@ -73,7 +73,7 @@ testinput_arrays_only = {
         {"bar": [testinput["testfoo"][0]["bar"][0]]},
         {"bar": []},
     ],
-    "testbar": {"hey there!": {"bar": [testinput["testbar"]["hey there!"]["bar"][0]]}},
+    "testbar": {"hey there": {"bar": [testinput["testbar"]["hey there"]["bar"][0]]}},
     "testlazysequence": [
         (testinput["testlazysequence"][0][1],),
         (testinput["testlazysequence"][1][1],),
@@ -166,7 +166,7 @@ def test_create_jacobian_schema():
         jac_testinput[var] = [
             "testdiffarr",
             "testfoo.[0].bar.[0]",
-            "testbar.{hey there!}.bar.[0]",
+            "testbar.{hey there}.bar.[0]",
         ]
         InputSchema.model_validate(jac_testinput)
 
@@ -184,20 +184,49 @@ def test_create_jacobian_schema():
             jac_testinput[var] = ["moo"]
             InputSchema.model_validate(jac_testinput)
 
-    OutputSchema.model_validate(testoutput)
+    jac_testinput = valid_jac_input.copy()
+    # Test LookupError (from IndexError) (only raised on jac_inputs)
+    with pytest.raises(ValidationError):
+        jac_testinput["jac_inputs"] = ["testfoo.[2].bar.[0]"]
+        InputSchema.model_validate(jac_testinput)
+
+    with pytest.raises(ValidationError):
+        jac_testinput["jac_inputs"] = ["testlazysequence.[0].[2]"]
+        InputSchema.model_validate(jac_testinput)
+
+    # Test LookupError (from KeyError) (only raised on jac_inputs)
+    with pytest.raises(ValidationError):
+        jac_testinput["jac_inputs"] = ["testbar.{hey ther}.bar.[0]"]
+        InputSchema.model_validate(jac_testinput)
 
+    ctx = {
+        "input_keys": valid_jac_input["jac_inputs"],
+        "output_keys": valid_jac_input["jac_outputs"],
+    }
+
+    OutputSchema.model_validate(testoutput, context=ctx)
+
+    # Input instead of output
     with pytest.raises(ValidationError):
-        OutputSchema.model_validate(testinput)
+        OutputSchema.model_validate(testinput, context=ctx)
 
     # Extra keys
     with pytest.raises(ValidationError):
-        OutputSchema.model_validate({**testoutput, "foo": 1})
+        OutputSchema.model_validate({**testoutput, "foo": 1}, context=ctx)
 
     # Extra keys (nested)
     with pytest.raises(ValidationError):
         out = testoutput.copy()
         out["testdiffarr"]["foo"] = 1
-        OutputSchema.model_validate(out)
+        OutputSchema.model_validate(out, context=ctx)
+
+    # Missing keys
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate({**testoutput, "testdiffarr": {}}, context=ctx)
+
+    # Flat dict
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(testoutput["testdiffarr"], context=ctx)
 
 
 def test_create_jvp_schema():
@@ -221,28 +250,91 @@ def test_create_jvp_schema():
         jvp_testinput[var] = [
             "testdiffarr",
             "testfoo.[0].bar.[0]",
+            "testbar.{hey there}.bar.[0]",
             "testlazysequence.[0].[1]",
         ]
+        if var == "jvp_inputs":
+            jvp_testinput["tangent_vector"] = {
+                **jvp_testinput["tangent_vector"],
+                "testfoo.[0].bar.[0]": testinput["testfoo"][0]["bar"][0],
+                "testbar.{hey there}.bar.[0]": testinput["testbar"]["hey there"]["bar"][
+                    0
+                ],
+                "testlazysequence.[0].[1]": testinput["testlazysequence"][0][1],
+            }
         InputSchema.model_validate(jvp_testinput)
 
         with pytest.raises(ValidationError):
             jvp_testinput[var] = []
+            if var == "jvp_inputs":
+                jvp_testinput["tangent_vector"] = {}
             InputSchema.model_validate(jvp_testinput)
 
         with pytest.raises(ValidationError):
             # non-differentiable array
             jvp_testinput[var] = ["testbaz"]
+            if var == "jvp_inputs":
+                jvp_testinput["tangent_vector"] = {"testbaz": testinput["testbaz"]}
             InputSchema.model_validate(jvp_testinput)
 
         with pytest.raises(ValidationError):
             # non-existing
             jvp_testinput[var] = ["moo"]
+            if var == "jvp_inputs":
+                jvp_testinput["tangent_vector"] = {"moo": testinput["testdiffarr"]}
             InputSchema.model_validate(jvp_testinput)
 
-    OutputSchema.model_validate(testoutput)
+    jvp_testinput = valid_jvp_input.copy()
+    # Test LookupError (from IndexError) (only raised on jvp_inputs)
+    with pytest.raises(ValidationError):
+        jvp_testinput["jvp_inputs"] = ["testfoo.[2].bar.[0]"]
+        jvp_testinput["tangent_vector"] = {
+            "testfoo.[2].bar.[0]": testinput["testdiffarr"]
+        }
+        InputSchema.model_validate(jvp_testinput)
+
+    with pytest.raises(ValidationError):
+        jvp_testinput["jvp_inputs"] = ["testlazysequence.[0].[2]"]
+        jvp_testinput["tangent_vector"] = {
+            "testlazysequence.[0].[2]": testinput["testdiffarr"]
+        }
+        InputSchema.model_validate(jvp_testinput)
+
+    # Test LookupError (from KeyError) (only raised on jvp_inputs)
+    with pytest.raises(ValidationError):
+        jvp_testinput["jvp_inputs"] = ["testbar.{hey ther}.bar.[0]"]
+        jvp_testinput["tangent_vector"] = {
+            "testbar.{hey ther}.bar.[0]": testinput["testbar"]["hey there"]["bar"][0]
+        }
+        InputSchema.model_validate(jvp_testinput)
 
+    # Mis-aligned tangent vector
+    jvp_testinput = valid_jvp_input.copy()
     with pytest.raises(ValidationError):
-        OutputSchema.model_validate(testinput)
+        jvp_testinput["tangent_vector"] = {**jvp_testinput["tangent_vector"], "foo": 1}
+        InputSchema.model_validate(jvp_testinput)
+
+    OutputSchema.model_validate(
+        testoutput, context={"output_keys": valid_jvp_input["jvp_outputs"]}
+    )
+
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(
+            testinput, context={"output_keys": valid_jvp_input["jvp_outputs"]}
+        )
+
+    # Extra keys
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(
+            {**testoutput, "foo": 1},
+            context={"output_keys": valid_jvp_input["jvp_outputs"]},
+        )
+
+    # Missing keys
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(
+            {}, context={"output_keys": valid_jvp_input["jvp_outputs"]}
+        )
 
 
 def test_create_vjp_schema():
@@ -269,26 +361,79 @@ def test_create_vjp_schema():
             "testfoo.[0].bar.[0]",
             "testlazysequence.[0].[1]",
         ]
+        if var == "vjp_outputs":
+            vjp_testinput["cotangent_vector"] = {
+                **vjp_testinput["cotangent_vector"],
+                "testfoo.[0].bar.[0]": 1,
+                "testlazysequence.[0].[1]": testinput["testlazysequence"][0][1],
+            }
         InputSchema.model_validate(vjp_testinput)
 
         with pytest.raises(ValidationError):
             vjp_testinput[var] = []
+            if var == "vjp_outputs":
+                vjp_testinput["cotangent_vector"] = {}
             InputSchema.model_validate(vjp_testinput)
 
         with pytest.raises(ValidationError):
             # non-differentiable array
             vjp_testinput[var] = ["testbaz"]
+            if var == "vjp_outputs":
+                vjp_testinput["cotangent_vector"] = {"testbaz": testinput["testbaz"]}
             InputSchema.model_validate(vjp_testinput)
 
         with pytest.raises(ValidationError):
             # non-existing
             vjp_testinput[var] = ["moo"]
+            if var == "vjp_outputs":
+                vjp_testinput["cotangent_vector"] = {"moo": 1}
             InputSchema.model_validate(vjp_testinput)
 
-    OutputSchema.model_validate(testoutput)
+    vjp_testinput = valid_vjp_input.copy()
+    # Test LookupError (from IndexError) (only raised on vjp_inputs)
+    with pytest.raises(ValidationError):
+        vjp_testinput["vjp_inputs"] = ["testfoo.[2].bar.[0]"]
+        InputSchema.model_validate(vjp_testinput)
+
+    with pytest.raises(ValidationError):
+        vjp_testinput["vjp_inputs"] = ["testlazysequence.[0].[2]"]
+        InputSchema.model_validate(vjp_testinput)
 
+    # Test LookupError (from KeyError) (only raised on vjp_inputs)
     with pytest.raises(ValidationError):
-        OutputSchema.model_validate(testinput)
+        vjp_testinput["vjp_inputs"] = ["testbar.{hey ther}.bar.[0]"]
+        InputSchema.model_validate(vjp_testinput)
+
+    # Mis-aligned cotangent vector
+    vjp_testinput = valid_vjp_input.copy()
+    with pytest.raises(ValidationError):
+        vjp_testinput["cotangent_vector"] = {
+            **vjp_testinput["cotangent_vector"],
+            "foo": 1,
+        }
+        InputSchema.model_validate(vjp_testinput)
+
+    OutputSchema.model_validate(
+        testoutput, context={"input_keys": valid_vjp_input["vjp_inputs"]}
+    )
+
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(
+            testinput, context={"input_keys": valid_vjp_input["vjp_inputs"]}
+        )
+
+    # Extra keys
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(
+            {**testoutput, "foo": 1},
+            context={"input_keys": valid_vjp_input["vjp_inputs"]},
+        )
+
+    # Missing keys
+    with pytest.raises(ValidationError):
+        OutputSchema.model_validate(
+            {}, context={"input_keys": valid_vjp_input["vjp_inputs"]}
+        )
 
 
 def test_untyped_container_schema_generation():