text: from typing import List, Sequence
  
  import jax
  import jax.numpy as jnp
  import matplotlib.pyplot as plt
  import numpy as np
  
  import treex as tx
  
  
  class MLP(tx.Module):
      layers: List[tx.Linear]
  
      def __init__(self, features: Sequence[int]):
  
          self.layers = [tx.Linear(dout) for dout in features]
  
      def __call__(self, x):
          for linear in self.layers[:-1]:
              x = jax.nn.relu(linear(x))
          return self.layers[-1](x)
  
  
  x = np.random.uniform(-1, 1, size=(100, 1))
  y = 1.4 * x ** 2 - 0.3 + np.random.normal(scale=0.1, size=(100, 1))
  
  model = MLP([12, 8, 1]).init(42, x)
  
  
  @jax.jit
  @jax.grad
  def loss_fn(model, x, y):
      preds = model(x)
      return jnp.mean((preds - y) ** 2)
  
  
  # in reality use optax
  def sdg(param, grad):
      return param - 0.01 * grad
  
  
  # training loop
  for step in range(10_000):
      grads = loss_fn(model, x, y)
      model = jax.tree_map(sdg, model, grads)
  
  model = model.eval()
  preds = model(x)
  
  plt.plot(x, y, "o")
  plt.plot(x, preds, "x")
  plt.show()
  