Format example code (ruff 0.2.0)

hzdr · Feb 3, 2024 · bccd624 · bccd624
1 parent 9c25cf0
commit bccd624
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 24 deletions.
diff --git a/examples/test_ag.py b/examples/test_ag.py
@@ -23,7 +23,7 @@ def pow2(x):
 
 def pow2_vjp(ans, x):
     # correct for x scalar or (n,), use this in production code
-    return lambda v: v * 2*x
+    return lambda v: v * 2 * x
 
 
 def pow2_vjp_with_jac(ans, x):
@@ -32,10 +32,10 @@ def pow2_vjp_with_jac(ans, x):
     """
     # jacobian() works for scalar and 1d array input, diag() doesn't
     if x.shape == ():
-        return lambda v: v * 2*x
+        return lambda v: v * 2 * x
     else:
         ##jac = jacobian(lambda x: anp.power(x,2))(x)
-        jac = anp.diag(2*x)
+        jac = anp.diag(2 * x)
         return lambda v: anp.dot(v, jac)
 
 
@@ -67,7 +67,7 @@ def mysum_vjp(ans, x):
 
 
 def func(x):
-    return anp.sum(anp.power(anp.sin(x),2))
+    return anp.sum(anp.power(anp.sin(x), 2))
 
 
 def func_with_vjp(x):
@@ -79,35 +79,35 @@ def test():
     # df/dx : R -> R
     assert anp.allclose(grad(anp.sin)(1.234), anp.cos(1.234))
 
-    x = rand(10)*5 - 5
+    x = rand(10) * 5 - 5
     assert anp.allclose(jacobian(anp.sin)(x), anp.diag(anp.cos(x)))
 
     # elementwise_grad(f) : R^n -> R^n (of f: R^n -> R^n), returns the column sum of
     # the Jacobian
     assert anp.allclose(elementwise_grad(anp.sin)(x), anp.cos(x))
     assert anp.allclose(jacobian(anp.sin)(x).sum(axis=0), anp.cos(x))
 
-
     defvjp(mysin, mysin_vjp)
     defvjp(mysum, mysum_vjp)
     for p2_jvp in [pow2_vjp, pow2_vjp_with_jac]:
         defvjp(pow2, p2_jvp)
 
-        assert anp.allclose([func(xi)          for xi in x],
-                            [func_with_vjp(xi) for xi in x])
+        assert anp.allclose(
+            [func(xi) for xi in x], [func_with_vjp(xi) for xi in x]
+        )
 
-        assert anp.allclose(func(x),
-                            func_with_vjp(x))
+        assert anp.allclose(func(x), func_with_vjp(x))
 
-        assert anp.allclose([grad(func)(xi)          for xi in x],
-                            [grad(func_with_vjp)(xi) for xi in x])
+        assert anp.allclose(
+            [grad(func)(xi) for xi in x], [grad(func_with_vjp)(xi) for xi in x]
+        )
 
-        assert anp.allclose(elementwise_grad(func)(x),
-                            elementwise_grad(func_with_vjp)(x))
+        assert anp.allclose(
+            elementwise_grad(func)(x), elementwise_grad(func_with_vjp)(x)
+        )
 
-        assert anp.allclose(grad(func)(x),
-                            grad(func_with_vjp)(x))
+        assert anp.allclose(grad(func)(x), grad(func_with_vjp)(x))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     test()
diff --git a/examples/test_jax.py b/examples/test_jax.py
@@ -68,7 +68,7 @@ def mysin(x):
     much slower than np.sin() :-D
     """
     ##return jnp.sin(x)
-    return x - x ** 3 / 6 + x ** 5 / 120 - x ** 7 / 5040 + x ** 9 / 362880
+    return x - x**3 / 6 + x**5 / 120 - x**7 / 5040 + x**9 / 362880
 
 
 def mycos(x):
@@ -87,12 +87,7 @@ def mycos(x):
     slightly worse. Both grow beyond 1e-8 outside of ~ [-1,1].
     """
     return (
-        1
-        - x ** 2 / 2
-        + x ** 4 / 24
-        - x ** 6 / 720
-        + x ** 8 / 40320
-        - x ** 10 / 3628800
+        1 - x**2 / 2 + x**4 / 24 - x**6 / 720 + x**8 / 40320 - x**10 / 3628800
     )