<jupyter_start><jupyter_text><jupyter_code>import jax
import jax.numpy as np
import numpy as onp
import matplotlib.pyplot as plt
def func(x):
    return np.sin(x * 4 * np.pi)

grad_sine = jax.jit(jax.grad(func))
grad_grad_sine = jax.jit(jax.grad(grad_sine))
  
plt.rcParams['figure.figsize'] = [18, 6]
plt.figure()
plt.suptitle(r"$y=sin(4 \pi x)$", size=16)
# ---
ct_x = onp.linspace(0.0, 1.0, 1024)
ct_y = func(ct_x)
ax = plt.subplot(1, 3, 1)
ax.set_xlim(0, 1)
ax.set_yticks(onp.linspace(-2 ,2, 5))
ax.set_ylim(-2, 2)
ax.grid(color='k', linestyle='--')
ax.set_title("y")
ax.plot(ct_x, ct_y, 'k')
# ---
ct_x = onp.linspace(0.0, 1.0, 1024)
ct_y = np.array([grad_sine(xi) for xi in ct_x])
ax = plt.subplot(1, 3, 2)
ax.set_xlim(0, 1)
EXTENT = 8 * np.pi
ax.set_yticks(onp.linspace(-EXTENT, EXTENT, 5))
ax.set_ylim(-EXTENT, EXTENT)
ax.grid(color='k', linestyle='--')
ax.set_title(r"$dy/dx$")
ax.plot(ct_x, ct_y, 'k')
# ---
ct_x = onp.linspace(0.0, 1.0, 1024)
ct_y = np.array([grad_grad_sine(xi) for xi in ct_x])
ax = plt.subplot(1, 3, 3)
ax.set_xlim(0, 1)
EXTENT = 32 * np.pi * np.pi
ax.set_yticks(onp.linspace(-EXTENT, EXTENT, 5))
ax.set_ylim(-EXTENT, EXTENT)
ax.grid(color='k', linestyle='--')
ax.set_title(r"$d^2y/dx^2$")
ax.plot(ct_x, ct_y, 'k')
# ---
plt.show()

def func(x):
    return (x ** 3)

grad_sine = jax.jit(jax.grad(func))
grad_grad_sine = jax.jit(jax.grad(grad_sine))
  
plt.rcParams['figure.figsize'] = [18, 6]
plt.figure()
plt.suptitle(r"$y=x^3$", size=16)
# ---
ct_x = onp.linspace(0.0, 1.0, 1024)
ct_y = func(ct_x)
ax = plt.subplot(1, 3, 1)
ax.set_xlim(0, 1)
ax.grid(color='k', linestyle='--')
ax.set_title("y")
ax.plot(ct_x, ct_y, 'k')
# ---
ct_x = onp.linspace(0.0, 1.0, 1024)
ct_y = np.array([grad_sine(xi) for xi in ct_x])
ax = plt.subplot(1, 3, 2)
ax.set_xlim(0, 1)
ax.grid(color='k', linestyle='--')
ax.set_title(r"$dy/dx$")
ax.plot(ct_x, ct_y, 'k')
# ---
ct_x = onp.linspace(0.0, 1.0, 1024)
ct_y = np.array([grad_grad_sine(xi) for xi in ct_x])
ax = plt.subplot(1, 3, 3)
ax.set_xlim(0, 1)
ax.grid(color='k', linestyle='--')
ax.set_title(r"$d^2y/dx^2$")
ax.plot(ct_x, ct_y, 'k')
# ---
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.2331002331002331, 'fraction_numerical': 0.04941724941724942, 'mean_word_length': 4.460559796437659, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '313172', 'license': 'Unlicense', 'n_tokens_mistral': 1193, 'n_tokens_neox': 1127, 'n_words': 195, 'path': 'jax_basic_grad.ipynb', 'repo_name': 'ApoorvaJ/public-notebooks', 'size': '2145'}