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: Train this model on a dataset.
  Parameters
  ----------
  dataset: Dataset
      the Dataset to train on
  nb_epoch: int
      the number of epochs to train for
  deterministic: bool
      if True, the samples are processed in order.  If False, a different random
      order is used for each epoch.
  loss: function
      a function of the form f(outputs, labels, weights) that computes the loss
      for each batch.  If None (the default), the model's standard loss function
      is used.
  callbacks: function or list of functions
      one or more functions of the form f(model, step) that will be invoked after
      every step.  This can be used to perform validation, logging, etc.
  all_losses: Optional[List[float]], optional (default None)
      If specified, all logged losses are appended into this list. Note that
      you can call `fit()` repeatedly with the same list and losses will
      continue to be appended.
  Returns
  -------
  The average loss over the most recent checkpoint interval
  Miscellanous Parameters Yet To Add
  ----------------------------------
  max_checkpoints_to_keep: int
      the maximum number of checkpoints to keep.  Older checkpoints are discarded.
  checkpoint_interval: int
      the frequency at which to write checkpoints, measured in training steps.
      Set this to 0 to disable automatic checkpointing.
  restore: bool
      if True, restore the model from the most recent checkpoint and continue training
      from there.  If False, retrain the model from scratch.
  variables: list of hk.Variable
      the variables to train.  If None (the default), all trainable variables in
      the model are used.
  Work in Progress
  ----------------
  [1] Integerate the optax losses, optimizers, schedulers with Deepchem
  [2] Support for saving & loading the model.
  [3] Adding support for output types (choosing only self._loss_outputs)
code: def fit(
      self,
      dataset: Dataset,
      nb_epochs: int = 10,
      deterministic: bool = False,
      loss: Optional[Union[Loss, LossFn]] = None,
      callbacks: Union[Callable, List[Callable]] = [],
      all_losses: Optional[List[float]] = None,
  ) -> float:
      """Train this model on a dataset.
      Parameters
      ----------
      dataset: Dataset
          the Dataset to train on
      nb_epoch: int
          the number of epochs to train for
      deterministic: bool
          if True, the samples are processed in order.  If False, a different random
          order is used for each epoch.
      loss: function
          a function of the form f(outputs, labels, weights) that computes the loss
          for each batch.  If None (the default), the model's standard loss function
          is used.
      callbacks: function or list of functions
          one or more functions of the form f(model, step) that will be invoked after
          every step.  This can be used to perform validation, logging, etc.
      all_losses: Optional[List[float]], optional (default None)
          If specified, all logged losses are appended into this list. Note that
          you can call `fit()` repeatedly with the same list and losses will
          continue to be appended.
      Returns
      -------
      The average loss over the most recent checkpoint interval
      Miscellanous Parameters Yet To Add
      ----------------------------------
      max_checkpoints_to_keep: int
          the maximum number of checkpoints to keep.  Older checkpoints are discarded.
      checkpoint_interval: int
          the frequency at which to write checkpoints, measured in training steps.
          Set this to 0 to disable automatic checkpointing.
      restore: bool
          if True, restore the model from the most recent checkpoint and continue training
          from there.  If False, retrain the model from scratch.
      variables: list of hk.Variable
          the variables to train.  If None (the default), all trainable variables in
          the model are used.
      Work in Progress
      ----------------
      [1] Integerate the optax losses, optimizers, schedulers with Deepchem
      [2] Support for saving & loading the model.
      [3] Adding support for output types (choosing only self._loss_outputs)
      """
      return self.fit_generator(
          self.default_generator(
              dataset, epochs=nb_epochs, deterministic=deterministic
          ),
          loss,
          callbacks,
          all_losses,
      )
  
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 3 for calling the JaxModel obj's underlying member method code with example usage: obj.fit(dataset, nb_epochs, deterministic, loss, callbacks, all_losses) and returns: self