class_docstr: This is a DeepChem model implemented by a Jax Model
  Here is a simple example of that uses JaxModel to train a
  Haiku (JAX Neural Network Library) based model on deepchem
  dataset.
  
  >>>
  >> def forward_model(x):
  >>   net = hk.nets.MLP([512, 256, 128, 1])
  >>   return net(x)
  >> def rms_loss(pred, tar, w):
  >>   return jnp.mean(optax.l2_loss(pred, tar))
  >> params_init, forward_fn = hk.transform(forward_model)
  >> rng = jax.random.PRNGKey(500)
  >> inputs, _, _, _ = next(iter(dataset.iterbatches(batch_size=256)))
  >> params = params_init(rng, inputs)
  >> j_m = JaxModel(forward_fn, params, rms_loss, 256, 0.001, 100)
  >> j_m.fit(train_dataset)
  
  All optimizations will be done using the optax library.
docstr: Parameters
  ----------
  generator: generator
      this should generate batches, each represented as a tuple of the form
      (inputs, labels, weights).
  transformers: List[dc.trans.Transformers]
      Transformers that the input data has been transformed by.  The output
      is passed through these transformers to undo the transformations.
  output_types: String or list of Strings
      If specified, all outputs of this type will be retrieved
      from the model. If output_types is specified, outputs must
      be None.
  
  Returns
  -------
      a NumPy array of the model produces a single output, or a list of arrays
      if it produces multiple outputs
code: def predict_on_generator(
      self,
      generator: Iterable[Tuple[Any, Any, Any]],
      transformers: List[Transformer] = [],
      output_types: Optional[OneOrMany[str]] = None,
  ) -> OneOrMany[np.ndarray]:
      """
      Parameters
      ----------
      generator: generator
          this should generate batches, each represented as a tuple of the form
          (inputs, labels, weights).
      transformers: List[dc.trans.Transformers]
          Transformers that the input data has been transformed by.  The output
          is passed through these transformers to undo the transformations.
      output_types: String or list of Strings
          If specified, all outputs of this type will be retrieved
          from the model. If output_types is specified, outputs must
          be None.
  
      Returns
      -------
          a NumPy array of the model produces a single output, or a list of arrays
          if it produces multiple outputs
      """
      return self._predict(generator, transformers, False, output_types)
  
imports: import numpy
  import time
  import logging
  from collections.abc import Sequence
  from deepchem.data import Dataset, NumpyDataset
  from deepchem.metrics import Metric
  from deepchem.models.models import Model
  from deepchem.models.losses import Loss
  from deepchem.models.optimizers import Optimizer, Adam
  from deepchem.utils.evaluate import GeneratorEvaluator
  from deepchem.trans.transformers import Transformer, undo_transforms
  from typing import Any, Callable, Iterable, List, Optional, Tuple, Union, Sequence
  from deepchem.utils.typing import LossFn, OneOrMany
  import jax.numpy
  import jax
  import haiku
  import optax
  import warnings
label_desc: Use image node_id 6 for calling the JaxModel obj's underlying member method code with example usage: obj.predict_on_generator(generator, transformers, output_types) and returns: self