code: def test_overfit_subclass_model():
      """Test fitting a JaxModel defined by subclassing Module."""
      n_data_points = 10
      n_features = 2
      np.random.seed(1234)
      X = np.random.rand(n_data_points, n_features)
      y = (X[:, 0] > X[:, 1]).astype(np.float32)
      dataset = dc.data.NumpyDataset(X, np.expand_dims(y, axis=1))
  
      class Encoder(hk.Module):
          def __init__(self, output_size: int = 1):
              super().__init__()
              self._network = hk.nets.MLP([512, 256, 128, output_size])
  
          def __call__(self, x: jnp.ndarray):
              x = self._network(x)
              return x, jax.nn.sigmoid(x)
  
      # Model Initilisation
      params_init, forward_fn = hk.transform(
          lambda x: Encoder()(x)
      )  # 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 = lambda pred, tar, w: jnp.mean(  # noqa: E731
          optax.sigmoid_binary_cross_entropy(pred[0], tar)
      )  # noqa
  
      # JaxModel Working
      j_m = JaxModel(
          forward_fn,
          params,
          criterion,
          output_types=["loss", "prediction"],
          batch_size=100,
          learning_rate=0.001,
          log_frequency=2,
      )
      j_m.fit(dataset, nb_epochs=1000)
      prediction = np.squeeze(j_m.predict_on_batch(X))
      assert np.array_equal(y, np.round(prediction))
      metric = dc.metrics.Metric(dc.metrics.roc_auc_score)
      scores = j_m.evaluate(dataset, [metric])
      assert scores[metric.name] > 0.9
  
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 4 for calling a global function with example usage: test_overfit_subclass_model() and returns: x,  jax