Repository: google/jax
Issue: jax.hessian uses positional argument `holomorphic` where it should be `has_aux` 

EDIT: Fixed a Typo.

Hi, this is a bit of a minor issue, but the jax source implementation of `hessian` has a slight mixup (or misses) in the function arguments. It uses `holomorphic` in place of `has_aux`, which is slightly confusing.

As shown below (see https://jax.readthedocs.io/en/latest/_modules/jax/_src/api.html):
```python
def hessian(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
            holomorphic: bool = False) -> Callable:
  """Doc ... """
  return jacfwd(jacrev(fun, argnums, holomorphic), argnums, holomorphic)
```

However, the argument `holomorphic` is used in the positional arg of `has_aux` for both jacfwd and jacrev:
```python
def jacfwd(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
           has_aux: bool = False, holomorphic: bool = False) -> Callable:
    ...
```

Current behaviour:
```python
import jax


def f(x):
    return jax.numpy.sin(x)


def f_aux(x):
    return f(x), x


print(jax.hessian(f)(10.))
# 0.5440211
print(jax.hessian(f_aux, holomorphic=True)(10.))
# (DeviceArray(0.5440211, dtype=float32, weak_type=True), 10.0)
print(jax.hessian(f, holomorphic=True)(10.))
# TypeError: expected function with aux output to return a two-element tuple, but got type ...
```

Should be:
```python
import jax


def f(x):
    return jax.numpy.sin(x)


def f_aux(x):
    return f(x), x


print(jax.hessian(f)(10.))
# 0.5440211
print(jax.hessian(f_aux, has_aux=True)(10.))
# (DeviceArray(0.5440211, dtype=float32, weak_type=True), 10.0)
print(jax.hessian(f, has_aux=True)(10.))
# TypeError: expected function with aux output to return a two-element tuple, but got type ...
```


Fix:
```python
def hessian(
        fun: Callable,
        argnums: Union[int, Sequence[int]] = 0,
        has_aux: bool = False,
        holomorphic: bool = False
) -> Callable:
    """Doc ..."""
    return jacrev(jacfwd(fun, argnums, has_aux=has_aux, holomorphic=holomorphic),
                  argnums, has_aux=has_aux, holomorphic=holomorphic)
```


Changed files:
tests/api_test.py, jax/_src/api.py

diff --git a/jax/_src/api.py b/jax/_src/api.py
index 694345e11..f00831b3b 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1158,7 +1158,7 @@ _check_output_dtype_jacrev = partial(_check_output_dtype_revderiv, "jacrev")
 
 
 def hessian(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
-            holomorphic: bool = False) -> Callable:
+            has_aux: bool = False, holomorphic: bool = False) -> Callable:
   """Hessian of ``fun`` as a dense array.
 
   Args:
@@ -1168,6 +1168,9 @@ def hessian(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
       containers thereof.
     argnums: Optional, integer or sequence of integers. Specifies which
       positional argument(s) to differentiate with respect to (default ``0``).
+    has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the
+      first element is considered the output of the mathematical function to be
+      differentiated and the second element is auxiliary data. Default False.
     holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be
       holomorphic. Default False.
 
@@ -1218,7 +1221,8 @@ def hessian(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
   ``(out1, out2, ..., in1, in2, ..., in1, in2, ...)``. To flatten pytrees into
   1D vectors, consider using :py:func:`jax.flatten_util.flatten_pytree`.
   """
-  return jacfwd(jacrev(fun, argnums, holomorphic), argnums, holomorphic)
+  return jacfwd(jacrev(fun, argnums, has_aux=has_aux, holomorphic=holomorphic),
+                argnums, has_aux=has_aux, holomorphic=holomorphic)
 
 def _std_basis(pytree):
   leaves, _ = tree_flatten(pytree)
diff --git a/tests/api_test.py b/tests/api_test.py
index a54a96161..8e4ef94e4 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1217,6 +1217,26 @@ class APITest(jtu.JaxTestCase):
     f = lambda x: jnp.dot(x, jnp.dot(A, x))
     assert np.allclose(hessian(f)(x), A + A.T)
 
+  @jtu.skip_on_devices("tpu")
+  def test_hessian_holomorphic(self):
+    R = self.rng().randn
+    A = R(4, 4)
+    x = R(4) * (1 + 2j)
+
+    f = lambda x: jnp.dot(x, jnp.dot(A, x))
+    assert np.allclose(hessian(f, holomorphic=True)(x), A + A.T)
+
+  @jtu.skip_on_devices("tpu")
+  def test_hessian_aux(self):
+    R = self.rng().randn
+    A = R(4, 4)
+    x = R(4)
+
+    f = lambda x: (jnp.dot(x, jnp.dot(A, x)), x)
+    h, aux = hessian(f, has_aux=True)(x)
+    assert np.allclose(h, A + A.T)
+    assert np.allclose(aux, x)
+
   def test_std_basis(self):
     basis = api._std_basis(jnp.zeros(3))
     assert getattr(basis, "shape", None) == (3, 3)