code: def test_pure_jax_model():
      """
      Here we train a fully NN model made purely in Jax.
      The model is taken from Jax Tutorial https://jax.readthedocs.io/en/latest/notebooks/neural_network_with_tfds_data.html
      """
      n_data_points = 50
      n_features = 1
      np.random.seed(1234)
      X = np.random.rand(n_data_points, n_features)
      y = X * X + X + 1
      dataset = dc.data.NumpyDataset(X, y)
  
      # Initialize the weights with random values
      def random_layer_params(m, n, key, scale=1e-2):
          w_key, b_key = random.split(key)
          return scale * random.normal(
              w_key, (m, n)
          ), scale * random.normal(b_key, (n,))
  
      def init_network_params(sizes, key):
          keys = random.split(key, len(sizes))
          return [
              random_layer_params(m, n, k)
              for m, n, k in zip(sizes[:-1], sizes[1:], keys)
          ]
  
      layer_sizes = [1, 256, 128, 1]
      params = init_network_params(layer_sizes, random.PRNGKey(0))
  
      # Forward function which takes the params
      def forward_fn(params, rng, x):
          for i, weights in enumerate(params[:-1]):
              w, b = weights
              x = jnp.dot(x, w) + b
              x = jax.nn.relu(x)
  
          final_w, final_b = params[-1]
          output = jnp.dot(x, final_w) + final_b
          return output
  
      def rms_loss(pred, tar, w):
          return jnp.mean(optax.l2_loss(pred, tar))
  
      # Loss Function
      criterion = rms_loss
  
      # JaxModel Working
      j_m = JaxModel(
          forward_fn,
          params,
          criterion,
          batch_size=100,
          learning_rate=0.001,
          log_frequency=2,
      )
      j_m.fit(dataset, nb_epochs=1000)
      metric = dc.metrics.Metric(
          dc.metrics.mean_absolute_error, mode="regression"
      )
      scores = j_m.evaluate(dataset, [metric])
      assert scores[metric.name] < 0.5
  
file_path: deepchem/deepchem/models/jax_models/tests/test_jax_model.py
imports: import pytest
  from deepchem.models.tests.test_graph_models import get_dataset
  import deepchem
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: test_pure_jax_model() and returns: unknown,  unknown, unknown, output, jnp