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: Create a generator that iterates batches for a dataset.
  Subclasses may override this method to customize how model inputs are
  generated from the data.
  Parameters
  ----------
  dataset: Dataset
      the data to iterate
  epochs: int
      the number of times to iterate over the full dataset
  mode: str
      allowed values are 'fit' (called during training), 'predict' (called
      during prediction), and 'uncertainty' (called during uncertainty
      prediction)
  deterministic: bool
      whether to iterate over the dataset in order, or randomly shuffle the
      data for each epoch
  pad_batches: bool
      whether to pad each batch up to this model's preferred batch size
  Returns
  -------
  a generator that iterates batches, each represented as a tuple of lists:
  ([inputs], [outputs], [weights])
code: def default_generator(
      self,
      dataset: Dataset,
      epochs: int = 1,
      mode: str = "fit",
      deterministic: bool = True,
      pad_batches: bool = True,
  ) -> Iterable[Tuple[List, List, List]]:
      """Create a generator that iterates batches for a dataset.
      Subclasses may override this method to customize how model inputs are
      generated from the data.
      Parameters
      ----------
      dataset: Dataset
          the data to iterate
      epochs: int
          the number of times to iterate over the full dataset
      mode: str
          allowed values are 'fit' (called during training), 'predict' (called
          during prediction), and 'uncertainty' (called during uncertainty
          prediction)
      deterministic: bool
          whether to iterate over the dataset in order, or randomly shuffle the
          data for each epoch
      pad_batches: bool
          whether to pad each batch up to this model's preferred batch size
      Returns
      -------
      a generator that iterates batches, each represented as a tuple of lists:
      ([inputs], [outputs], [weights])
      """
  
      for epoch in range(epochs):
          for X_b, y_b, w_b, _ in dataset.iterbatches(
              batch_size=self.batch_size,
              deterministic=deterministic,
              pad_batches=pad_batches,
          ):
              yield ([X_b], [y_b], [w_b])
  
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 16 for calling the JaxModel obj's underlying member method code with example usage: obj.default_generator(dataset, epochs, mode, deterministic, pad_batches) without return types