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.
code: def fit_generator(
      self,
      generator: Iterable[Tuple[Any, Any, Any]],
      loss: Optional[Union[Loss, LossFn]] = None,
      callbacks: Union[Callable, List[Callable]] = [],
      all_losses: Optional[List[float]] = None,
  ) -> float:
      if not isinstance(callbacks, SequenceCollection):
          callbacks = [callbacks]
      self._ensure_built()
      avg_loss = 0.0
      last_avg_loss = 0.0
      averaged_batches = 0
      if loss is None:
          loss = self._loss_fn
      model_loss_fn = self._create_gradient_fn(
          self.forward_fn, self._loss_outputs, loss
      )
      grad_update = self._create_update_fn(
          self.optimizer, model_loss_fn
      )
  
      params, opt_state = self._get_trainable_params()
      rng = self.rng
      time1 = time.time()
  
      # Main training loop
  
      for batch in generator:
          inputs, labels, weights = self._prepare_batch(batch)
  
          if isinstance(inputs, list) and len(inputs) == 1:
              inputs = inputs[0]
  
          if isinstance(labels, list) and len(labels) == 1:
              labels = labels[0]
  
          if isinstance(weights, list) and len(weights) == 1:
              weights = weights[0]
  
          params, opt_state, batch_loss = grad_update(
              params, opt_state, inputs, labels, weights, rng=rng
          )
          rng, _ = jax.random.split(rng)
          avg_loss += jax.device_get(batch_loss)
          self._global_step += 1
          current_step = self._global_step
          averaged_batches += 1
          should_log = current_step % self.log_frequency == 0
  
          if should_log:
              avg_loss = float(avg_loss) / averaged_batches
              logger.info(
                  "Ending global_step %d: Average loss %g"
                  % (current_step, avg_loss)
              )
              if all_losses is not None:
                  all_losses.append(avg_loss)
              # Capture the last avg_loss in case of return since we're resetting to 0 now
              last_avg_loss = avg_loss
              avg_loss = 0.0
              averaged_batches = 0
          for c in callbacks:
              c(self, current_step)
  
      # Report final results.
      if averaged_batches > 0:
          avg_loss = float(avg_loss) / averaged_batches
          logger.info(
              "Ending global_step %d: Average loss %g"
              % (current_step, avg_loss)
          )
          if all_losses is not None:
              all_losses.append(avg_loss)
          last_avg_loss = avg_loss
  
      time2 = time.time()
      logger.info(
          "TIMING: model fitting took %0.3f s" % (time2 - time1)
      )
      self._set_trainable_params(params, opt_state)
      return last_avg_loss
  
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 4 for calling the JaxModel obj's underlying member method code with example usage: obj.fit_generator(generator, loss, callbacks, all_losses) and returns: last_avg_loss