added jacobian tests

Author: jpbrodrick89

tests/nested_tesseract/tesseract_api.py

@@ -81,6 +81,17 @@ def vector_jacobian_product(
     return out
 
 
+def jacobian(inputs: InputSchema, jac_inputs: set[str], jac_outputs: set[str]):
+    jac = {dy: {dx: [0.0, 0.0, 0.0] for dx in jac_inputs} for dy in jac_outputs}
+
+    if "scalars.a" in jac_inputs and "scalars.a" in jac_outputs:
+        jac["scalars.a"]["scalars.a"] = 10.0
+    if "vectors.v" in jac_inputs and "vectors.v" in jac_outputs:
+        jac["vectors.v"]["vectors.v"] = [[10.0, 0, 0], [0, 10.0, 0], [0, 0, 10.0]]
+
+    return jac
+
+
 def abstract_eval(abstract_inputs):
     """Calculate output shape of apply from the shape of its inputs."""
     return {

tests/test_endtoend.py

@@ -36,7 +36,7 @@ def _assert_pytree_isequal(a, b, rtol=None, atol=None):
             else:
                 assert a_elem == b_elem, f"Values are different: {a_elem} != {b_elem}"
         except AssertionError as e:
-            failures.append(a_path, str(e))
+            failures.append((a_path, str(e)))
 
     if failures:
         msg = "\n".join(f"Path: {path}, Error: {error}" for path, error in failures)
@@ -148,6 +148,39 @@ def f(x, y):
     _assert_pytree_isequal(vjp, vjp_raw)
 
 
+@pytest.mark.parametrize("use_jit", [True, False])
+@pytest.mark.parametrize("jacfun", [jax.jacfwd, jax.jacrev])
+def test_univariate_tesseract_jacobian(
+    served_univariate_tesseract_raw, use_jit, jacfun
+):
+    rosenbrock_tess = Tesseract(served_univariate_tesseract_raw)
+
+    # make things callable without keyword args
+    def f(x, y):
+        return apply_tesseract(rosenbrock_tess, inputs=dict(x=x, y=y))["result"]
+
+    rosenbrock_raw = rosenbrock_impl
+    if use_jit:
+        f = jax.jit(f)
+        rosenbrock_raw = jax.jit(rosenbrock_raw)
+
+    x, y = np.array(0.0), np.array(0.0)
+    jac = jacfun(f, argnums=(0, 1))(x, y)
+
+    # Test against Tesseract client
+    jac_ref = rosenbrock_tess.jacobian(
+        inputs=dict(x=x, y=y), jac_inputs=["x", "y"], jac_outputs=["result"]
+    )
+
+    # Convert from nested dict to nested tuplw
+    jac_ref = tuple((jac_ref["result"]["x"], jac_ref["result"]["y"]))
+    _assert_pytree_isequal(jac, jac_ref)
+
+    # Test against direct implementation
+    jac_raw = jacfun(rosenbrock_raw, argnums=(0, 1))(x, y)
+    _assert_pytree_isequal(jac, jac_raw)
+
+
 @pytest.mark.parametrize("use_jit", [True, False])
 def test_nested_tesseract_apply(served_nested_tesseract_raw, use_jit):
     nested_tess = Tesseract(served_nested_tesseract_raw)
@@ -286,6 +319,55 @@ def f(a, v):
     _assert_pytree_isequal(vjp, vjp_ref)
 
 
+@pytest.mark.parametrize("use_jit", [True, False])
+@pytest.mark.parametrize("jacfun", [jax.jacfwd, jax.jacrev])
+def test_nested_tesseract_jacobian(served_nested_tesseract_raw, use_jit, jacfun):
+    nested_tess = Tesseract(served_nested_tesseract_raw)
+    a, b = np.array(1.0, dtype="float32"), np.array(2.0, dtype="float32")
+    v, w = (
+        np.array([1.0, 2.0, 3.0], dtype="float32"),
+        np.array([5.0, 7.0, 9.0], dtype="float32"),
+    )
+
+    def f(a, v):
+        return apply_tesseract(
+            nested_tess,
+            inputs=dict(
+                scalars={"a": a, "b": b},
+                vectors={"v": v, "w": w},
+                other_stuff={"s": "hey!", "i": 1234, "f": 2.718},
+            ),
+        )
+
+    if use_jit:
+        f = jax.jit(f)
+
+    jac = jacfun(f, argnums=(0, 1))(a, v)
+
+    jac_ref = nested_tess.jacobian(
+        inputs=dict(
+            scalars={"a": a, "b": b},
+            vectors={"v": v, "w": w},
+            other_stuff={"s": "hey!", "i": 1234, "f": 2.718},
+        ),
+        jac_inputs=["scalars.a", "vectors.v"],
+        jac_outputs=["scalars.a", "vectors.v"],
+    )
+    # JAX returns a 2-layered nested dict containing tuples of arrays
+    # we need to flatten it to match the Tesseract output (2 layered nested dict of arrays)
+    jac = {
+        "scalars.a": {
+            "scalars.a": jac["scalars"]["a"][0],
+            "vectors.v": jac["scalars"]["a"][1],
+        },
+        "vectors.v": {
+            "scalars.a": jac["vectors"]["v"][0],
+            "vectors.v": jac["vectors"]["v"][1],
+        },
+    }
+    _assert_pytree_isequal(jac, jac_ref)
+
+
 @pytest.mark.parametrize("use_jit", [True, False])
 def test_partial_differentiation(served_univariate_tesseract_raw, use_jit):
     """Test that differentiation works correctly in cases where some inputs are constants."""