more to torch.

Author: v0lta

README.md

@@ -43,8 +43,8 @@ Once more the problem is illustrated below:
 
 $$ \min_{\mathbf{x}} \mathbf{x} \cdot \mathbf{x} + \cos(2 \pi x_0 ) + \sin(2 \pi x_1)  + 0.5 \cdot \text{relu}(x_0) + 10 \cdot \tanh( \|\mathbf{x} \| ),  \text{   with   } \mathbf{x_0} = (2.9, -2.9) .$$
 
-The function is already defined in  `src/optimize_2d_momentum_bumpy_jax.py`. We dont have to find the gradient by hand!
-Use `jax.grad` [(jax-documentation)](https://jax.readthedocs.io/en/latest/_autosummary/jax.grad.html) to compute the gradient automatically. Use the result to find the minimum using momentum.  
+The function is already defined in  `src/optimize_2d_momentum_bumpy_torch.py`. We dont have to find the gradient by hand!
+Use `torch.func.grad` [(torch-documentation)](https://pytorch.org/docs/stable/generated/torch.func.grad.html) to compute the gradient automatically. Use the result to find the minimum using momentum.  
 
 While coding use `nox -s test`, `nox -s lint`, and `nox -s typing` to check your code.
 Autoformatting help is available via `nox -s format`.

noxfile.py

@@ -1,4 +1,5 @@
 """This module implements our CI function calls."""
+
 import nox
 
 

requirements.txt

@@ -1,4 +1,4 @@
 nox
 numpy
 matplotlib
-jax
+torch

src/optimize_1d.py

@@ -1,4 +1,5 @@
 """Implement gradient descent in 1d."""
+
 import matplotlib.pyplot as plt
 import numpy as np
 

src/optimize_2d_momentum_bumpy_jax.py src/optimize_2d_momentum_bumpy_torch.pysrc/optimize_2d_momentum_bumpy_jax.py renamed to src/optimize_2d_momentum_bumpy_torch.py

@@ -1,13 +1,13 @@
 """Implement 2d gradient descent using jax."""
 
-import jax
-import jax.numpy as np
-import numpy as nnp
+import numpy as np
+import torch as th
+from torch.func import grad
 
 from util import write_movie
 
 
-def bumpy_function(pos: np.ndarray) -> np.ndarray:
+def bumpy_function(pos: th.Tensor) -> th.Tensor:
     """Return values from an even bumpier function.
 
     This even bumpier function is hard to optimize.
@@ -22,10 +22,10 @@ def bumpy_function(pos: np.ndarray) -> np.ndarray:
     return (
         pos[0] * pos[0]
         + pos[1] * pos[1]
-        + np.cos(pos[0] * 2 * np.pi)
-        + np.sin(pos[1] * 2 * np.pi)
-        + (pos[0] > 0).astype(pos.dtype) * 0.5
-        + np.tanh(np.sqrt(pos[0] ** 2 + pos[1] ** 2)) * 10
+        + th.cos(pos[0] * 2 * th.pi)
+        + th.sin(pos[1] * 2 * th.pi)
+        + (pos[0] > 0).type(pos.dtype) * 0.5
+        + th.tanh(th.sqrt(pos[0] ** 2 + pos[1] ** 2)) * 10
     )
 
 
@@ -36,10 +36,10 @@ def bumpy_function(pos: np.ndarray) -> np.ndarray:
     # TODO: use jax to find the gradient.
 
     nx, ny = (1001, 1001)
-    x = np.linspace(-3, 3, nx)
-    y = np.linspace(-3, 3, ny)
-    mx, my = np.meshgrid(x, y)
-    pos = np.stack((mx, my))
+    x = th.linspace(-3, 3, nx)
+    y = th.linspace(-3, 3, ny)
+    mx, my = th.meshgrid(x, y)
+    pos = th.stack((mx, my))
     mz = bumpy_function(pos)
 
     fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
@@ -51,7 +51,7 @@ def bumpy_function(pos: np.ndarray) -> np.ndarray:
     plt.contourf(mx, my, mz)
     plt.colorbar()
 
-    start_pos = np.array((2.9, -2.9))
+    start_pos = th.tensor((2.9, -2.9))
     step_size = 0.0  # TODO: Choose your step size.
     alpha = 0.0  # TODO: Choose your momentum term.
     step_total = 100
@@ -65,9 +65,9 @@ def bumpy_function(pos: np.ndarray) -> np.ndarray:
     plt.show()
 
     write_movie(
-        nnp.array(mx),
-        nnp.array(my),
-        nnp.array(mz),
+        np.array(mx),
+        np.array(my),
+        np.array(mz),
         pos_list,
         "writer_grad_bumpy_plot_jax",
     )

src/util.py

@@ -1,4 +1,5 @@
 """Code to export gradient descent sequences into a movie."""
+
 from typing import Optional
 
 import matplotlib.animation as manimation

tests/test_function.py

@@ -5,6 +5,7 @@
 "Difference Qutiones may sometimes be useful, too"
 Andreas Griewank, Andrea Walther - Evaluating Derivatives.
 """
+
 import sys
 from typing import Callable, Optional, Union