upload changes

app.py

@@ -79,9 +79,10 @@ losses = []
 value_and_grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
 
 # initialize plot
-fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+fig, axes = plt.subplots(1, 2, figsize=(12, 6))
 axes[0].set_xlim(0, 1)
 axes[0].set_ylim(0, 1)
+axes[0].set_aspect("equal")
 
 value, aux = loss_fn(params)
 (pbola_plot,) = axes[0].plot(x, parabola_fn(x), color="red")