# Univariate Rosenbrock function

## Context
Example that wraps the univariate [Rosenbrock function](https://en.wikipedia.org/wiki/Rosenbrock_function), which is a common test problem for optimization algorithms. Defines a Tesseract that has all optional endpoints implemented, including `apply`, `abstract_eval`, `jacobian`, `jacobian_vector_product`, and `vector_jacobian_product`.

```{seealso}
This example (and using it to perform optimization) is also part of an [Expert Showcase](https://si-tesseract.discourse.group/t/jax-based-rosenbrock-function-minimization/48) in the Tesseract Community Forum.
```

## Example Tesseract (examples/univariate)

### Core functionality --- schemas and `apply` function

This example uses a pure-Python implementation of the Rosenbrock function as the basis for all endpoints:

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: rosenbrock
:language: python
```

As such, the Tesseract has 2 differentiable scalar inputs (`x` and `y`) and a single output (the value of the Rosenbrock function at those inputs). The parameters `a` and `b` are treated as non-differentiable constants.

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: InputSchema
:language: python
```

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: OutputSchema
:language: python
```

This makes it straightforward to write the `apply` function, which simply unpacks the inputs and calls the `rosenbrock` function with them:

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: apply
:language: python
```

### Jacobian endpoint

For the Jacobian, we exploit the fact that the `rosenbrock` function is traceable by [JAX](https://github.com/jax-ml/jax). We can therefore use [`jax.jacrev`](https://docs.jax.dev/en/latest/_autosummary/jax.jacrev.html) to compute the Jacobian of the function with respect to its inputs:

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: jacobian
:language: python
```

### Other AD endpoints

We define the JVP (Jacobian-vector product) and VJP (vector-Jacobian product) endpoints by summing over rows / columns of the Jacobian matrix. That is, we call `jacobian` under the hood, then multiply the resulting Jacobian matrix by the (tangent / cotangent) vector input.

```{warning}
Defining JVP and VJP operations through sums over the full Jacobian matrix is inefficient and negates the benefits of using JVP / VJP. These endpoints are provided for completeness, but in practice, you would typically use JAX's built-in JVP and VJP functions directly.
```

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: jacobian_vector_product
:language: python
```

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: vector_jacobian_product
:language: python
```

### Abstract evaluation

Some Tesseract clients (like [Tesseract-JAX](https://github.com/pasteurlabs/tesseract-jax)) require an abstract evaluation endpoint in order to pre-allocate memory for the inputs and outputs. This is a simple function that returns the shapes of the outputs based on the shapes of the inputs. In this case, the output is always a scalar, so we return an empty shape tuple:

```{literalinclude} ../../../../examples/univariate/tesseract_api.py
:pyobject: abstract_eval
:language: python
```
