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: Predict outputs for data provided by a generator.
  This is the private implementation of prediction.  Do not
  call it directly. Instead call one of the public prediction
  methods.
  
  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.
  uncertainty: bool
      specifies whether this is being called as part of estimating uncertainty.
      If True, it sets the training flag so that dropout will be enabled, and
      returns the values of the uncertainty outputs.
  other_output_types: list, optional
      Provides a list of other output_types (strings) to predict from model.
  
  Returns
  -------
  A NumpyArray if the model produces a single output, or a list of arrays otherwise.
code: def _predict(
      self,
      generator: Iterable[Tuple[Any, Any, Any]],
      transformers: List[Transformer],
      uncertainty: bool,
      other_output_types: Optional[OneOrMany[str]],
  ):
      """
      Predict outputs for data provided by a generator.
      This is the private implementation of prediction.  Do not
      call it directly. Instead call one of the public prediction
      methods.
  
      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.
      uncertainty: bool
          specifies whether this is being called as part of estimating uncertainty.
          If True, it sets the training flag so that dropout will be enabled, and
          returns the values of the uncertainty outputs.
      other_output_types: list, optional
          Provides a list of other output_types (strings) to predict from model.
  
      Returns
      -------
      A NumpyArray if the model produces a single output, or a list of arrays otherwise.
      """
      results: Optional[List[List[np.ndarray]]] = None
      variances: Optional[List[List[np.ndarray]]] = None
      if uncertainty and (other_output_types is not None):
          raise ValueError(
              "This model cannot compute uncertainties and other output types simultaneously. Please invoke one at a time."
          )
      if uncertainty:
          if (
              self._variance_outputs is None
              or len(self._variance_outputs) == 0
          ):
              raise ValueError(
                  "This model cannot compute uncertainties"
              )
          if len(self._variance_outputs) != len(
              self._prediction_outputs
          ):
              raise ValueError(
                  "The number of variances must exactly match the number of outputs"
              )
      if other_output_types:
          if (
              self._other_outputs is None
              or len(self._other_outputs) == 0
          ):
              raise ValueError(
                  "This model cannot compute other outputs since no other output_types were specified."
              )
      self._ensure_built()
      eval_fn = self._create_eval_fn(self.forward_fn, self.params)
      rng = self.rng
  
      for batch in generator:
          inputs, _, _ = self._prepare_batch(batch)
  
          if isinstance(inputs, list) and len(inputs) == 1:
              inputs = inputs[0]
  
          output_values = eval_fn(inputs, rng)
          if isinstance(output_values, jnp.ndarray):
              output_values = [output_values]
          output_values = [jax.device_get(t) for t in output_values]
  
          # Apply tranformers and record results.
          if uncertainty:
              var = [output_values[i] for i in self._variance_outputs]
              if variances is None:
                  variances = [var]
              else:
                  for i, t in enumerate(var):
                      variances[i].append(t)
  
          access_values = []
          if other_output_types:
              access_values += self._other_outputs
          elif self._prediction_outputs is not None:
              access_values += self._prediction_outputs
  
          if len(access_values) > 0:
              output_values = [output_values[i] for i in access_values]
  
          if len(transformers) > 0:
              if len(output_values) > 1:
                  raise ValueError(
                      "predict() does not support Transformers for models with multiple outputs."
                  )
              elif len(output_values) == 1:
                  output_values = [
                      undo_transforms(output_values[0], transformers)
                  ]
  
          if results is None:
              results = [[] for i in range(len(output_values))]
          for i, t in enumerate(output_values):
              results[i].append(t)
  
      # Concatenate arrays to create the final results.
      final_results = []
      final_variances = []
      if results is not None:
          for r in results:
              final_results.append(np.concatenate(r, axis=0))
      if uncertainty and variances is not None:
          for v in variances:
              final_variances.append(np.concatenate(v, axis=0))
          return zip(final_results, final_variances)
      if len(final_results) == 1:
          return final_results[0]
      else:
          return final_results
  
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 5 for calling the JaxModel obj's underlying member method code with example usage: obj._predict(generator, transformers, uncertainty, other_output_types) and returns: zip, final_results, final_results