code: def test_overfit_sequential_model():
      """Test fitting a JaxModel defined by subclassing Module."""
      n_data_points = 10
      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)
  
      def forward_fn(x):
          mlp = hk.Sequential(
              [
                  hk.Linear(300),
                  jax.nn.relu,
                  hk.Linear(100),
                  jax.nn.relu,
                  hk.Linear(1),
              ]
          )
          return mlp(x)
  
      def rms_loss(pred, tar, w):
          return jnp.mean(optax.l2_loss(pred, tar))
  
      # Model Initilisation
      params_init, forward_fn = hk.transform(forward_fn)  # noqa
      rng = jax.random.PRNGKey(500)
      inputs, _, _, _ = next(iter(dataset.iterbatches(batch_size=100)))
  
      modified_inputs = jnp.array(
          [
              x.astype(np.float32) if x.dtype == np.float64 else x
              for x in inputs
          ]
      )
      params = params_init(rng, modified_inputs)
  
      # 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 5 for calling a global function with example usage: test_overfit_sequential_model() and returns: mlp, jnp