doc: Python API for PyTorch example (#24)

#### Relevant issue or PR
Closes #6.

#### Description of changes
Context manager for `pytorch_optimization` example.

#### Testing done
CI + manual.

#### License

- [x] By submitting this pull request, I confirm that my contribution is
made under the terms of the [Apache 2.0
license](https://pasteurlabs.github.io/tesseract/LICENSE).
- [x] I sign the Developer Certificate of Origin below by adding my name
and email address to the `Signed-off-by` line.

<details>
<summary><b>Developer Certificate of Origin</b></summary>

```text
Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.
```

</details>

Signed-off-by: Matteo Salvador <matteo.salvador@simulation.science>

---------

Signed-off-by: Niklas Heim <niklas.heim@simulation.science>
Co-authored-by: Niklas Heim <niklas.heim@simulation.science>

Author: MatteoSalvador

examples/sciml/pytorch_optimization/optimization.py

@@ -1,16 +1,23 @@
 # Copyright 2025 Pasteur Labs. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+
 import argparse
 
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
-from tesseract_api import log_rosenbrock
-from tesseract_client import Client
+
+# disabling jit compiler for now
+import torch._dynamo
+
+torch._dynamo.config.disable = True
+
+from tesseract_api import log_rosenbrock  # noqa: E402
+
+from tesseract_core import Tesseract  # noqa: E402
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-p", "--port", help="Port of tesseract server", required=True)
 parser.add_argument("-n", "--niter", help="The number of iterations", default=300)
 parser.add_argument(
     "-l",
@@ -31,8 +38,6 @@
 
 args = vars(parser.parse_args())
 
-port = int(args["port"])
-
 torch.manual_seed(int(args["seed"]))
 
 # initiate starting x, y
@@ -53,9 +58,6 @@
     "b": b,
 }
 
-# Initialize Tesseract client
-client = Client(host="127.0.0.1", port=port)
-
 # setup optimizer: either Adam or SGD (default)
 assert args["opt"] in ["SGD", "Adam"], "Only Adam or SGD optimizer supported."
 opt = getattr(torch.optim, args["opt"])((x0, y0), lr=float(args["lr"]))
@@ -64,46 +66,41 @@
 losses = []
 x = []
 y = []
-for _i in range(int(args["niter"])):
-    # zero out the gradients
-    opt.zero_grad()
-
-    # get loss
-    output = client.request("apply", method="POST", payload={"inputs": inputs})
-    loss = output["loss"]["data"]["buffer"]
-
-    # append values
-    losses.append(loss)
-    x.append(inputs["x"])
-    y.append(inputs["y"])
-
-    # get gradients
-    jacobian_response = client.request(
-        "jacobian",
-        method="POST",
-        payload={
-            "inputs": inputs,
-            "jac_inputs": jac_inputs,
-            "jac_outputs": ["loss"],
-        },
-    )
-
-    # insert gradients into pytorch tensors where optimizer expects them
-    if "x" in jac_inputs:
-        x0.grad = torch.as_tensor(
-            jacobian_response["loss"]["x"]["data"]["buffer"]
-        ).unsqueeze(0)
-    if "y" in jac_inputs:
-        y0.grad = torch.as_tensor(
-            jacobian_response["loss"]["y"]["data"]["buffer"]
-        ).unsqueeze(0)
-
-    # make optimization step
-    opt.step()
-
-    # update inputs
-    inputs["x"] = x0.detach().item()
-    inputs["y"] = y0.detach().item()
+with Tesseract.from_image(image="pytorch_optimization") as pytorch_optimization:
+    for _i in range(int(args["niter"])):
+        # zero out the gradients
+        opt.zero_grad()
+
+        # get loss
+        output = pytorch_optimization.apply(inputs)
+        loss = output["loss"]
+
+        # append values
+        losses.append(loss)
+        x.append(inputs["x"])
+        y.append(inputs["y"])
+
+        # get gradients
+        jacobian_response = pytorch_optimization.jacobian(
+            inputs=inputs, jac_inputs=jac_inputs, jac_outputs=["loss"]
+        )
+
+        # insert gradients into pytorch tensors where optimizer expects them
+        if "x" in jac_inputs:
+            x0.grad = (
+                torch.as_tensor(jacobian_response["loss"]["x"]).float().unsqueeze(0)
+            )
+        if "y" in jac_inputs:
+            y0.grad = (
+                torch.as_tensor(jacobian_response["loss"]["y"]).float().unsqueeze(0)
+            )
+
+        # make optimization step
+        opt.step()
+
+        # update inputs
+        inputs["x"] = x0.detach().item()
+        inputs["y"] = y0.detach().item()
 
 # plotting
 X, Y = np.meshgrid(np.linspace(-3, 3, 100), np.linspace(-3, 3, 100), indexing="xy")