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 new JaxModel
  
  Parameters
  ----------
  model: hk.State or Function
      Any Jax based model that has a `apply` method for computing the network. Currently
      only haiku models are supported.
  params: hk.Params
      The parameter of the Jax based networks
  loss: dc.models.losses.Loss or function
      a Loss or function defining how to compute the training loss for each
      batch, as described above
  output_types: list of strings, optional (default None)
      the type of each output from the model, as described above
  batch_size: int, optional (default 100)
      default batch size for training and evaluating
  learning_rate: float or LearningRateSchedule, optional (default 0.001)
      the learning rate to use for fitting.  If optimizer is specified, this is
      ignored.
  optimizer: optax object
      For the time being, it is optax object
  rng: jax.random.PRNGKey, optional (default 1)
      A default global PRNG key to use for drawing random numbers.
  log_frequency: int, optional (default 100)
      The frequency at which to log data. Data is logged using
      `logging` by default.
  
  
  Miscellanous Parameters Yet To Add
  ----------------------------------
  model_dir: str, optional (default None)
      Will be added along with the save & load method
  tensorboard: bool, optional (default False)
      whether to log progress to TensorBoard during training
  wandb: bool, optional (default False)
      whether to log progress to Weights & Biases during training
  
  
  Work in Progress
  ----------------
  [1] Integrate the optax losses, optimizers, schedulers with Deepchem
  [2] Support for saving & loading the model.
code: def __init__(
      self,
      forward_fn: hk.State,
      params: hk.Params,
      loss: Optional[Union[Loss, LossFn]],
      output_types: Optional[List[str]] = None,
      batch_size: int = 100,
      learning_rate: float = 0.001,
      optimizer: Optional[
          Union[optax.GradientTransformation, Optimizer]
      ] = None,
      grad_fn: Callable = create_default_gradient_fn,
      update_fn: Callable = create_default_update_fn,
      eval_fn: Callable = create_default_eval_fn,
      rng=jax.random.PRNGKey(1),
      log_frequency: int = 100,
      **kwargs,
  ):
      """
      Create a new JaxModel
  
      Parameters
      ----------
      model: hk.State or Function
          Any Jax based model that has a `apply` method for computing the network. Currently
          only haiku models are supported.
      params: hk.Params
          The parameter of the Jax based networks
      loss: dc.models.losses.Loss or function
          a Loss or function defining how to compute the training loss for each
          batch, as described above
      output_types: list of strings, optional (default None)
          the type of each output from the model, as described above
      batch_size: int, optional (default 100)
          default batch size for training and evaluating
      learning_rate: float or LearningRateSchedule, optional (default 0.001)
          the learning rate to use for fitting.  If optimizer is specified, this is
          ignored.
      optimizer: optax object
          For the time being, it is optax object
      rng: jax.random.PRNGKey, optional (default 1)
          A default global PRNG key to use for drawing random numbers.
      log_frequency: int, optional (default 100)
          The frequency at which to log data. Data is logged using
          `logging` by default.
  
  
      Miscellanous Parameters Yet To Add
      ----------------------------------
      model_dir: str, optional (default None)
          Will be added along with the save & load method
      tensorboard: bool, optional (default False)
          whether to log progress to TensorBoard during training
      wandb: bool, optional (default False)
          whether to log progress to Weights & Biases during training
  
  
      Work in Progress
      ----------------
      [1] Integrate the optax losses, optimizers, schedulers with Deepchem
      [2] Support for saving & loading the model.
      """
      super(JaxModel, self).__init__(
          model=(forward_fn, params), **kwargs
      )
      warnings.warn(
          "JaxModel is still in active development and all features may not yet be implemented"
      )
      self._loss_fn = (
          loss  # lambda pred, tar: jnp.mean(optax.l2_loss(pred, tar))
      )
      self.batch_size = batch_size
      self.learning_rate = learning_rate
      if optimizer is None:
          optimizer = Adam(1e-3)
  
      if not isinstance(optimizer, optax.GradientTransformation):
          self.optimizer = optimizer._create_jax_optimizer()
      else:
          self.optimizer = optimizer
      self.forward_fn = forward_fn
      self.params = params
      self._built = False
      self.log_frequency = log_frequency
      self.rng = rng
      self._create_gradient_fn = grad_fn
      self._create_update_fn = update_fn
      self._create_eval_fn = eval_fn
  
      if output_types is None:
          self._prediction_outputs = None
          self._loss_outputs = None
          self._variance_outputs = None
          self._other_outputs = None
      else:
          self._prediction_outputs = []
          self._loss_outputs = []
          self._variance_outputs = []
          self._other_outputs = []
          for i, type in enumerate(output_types):
              if type == "prediction":
                  self._prediction_outputs.append(i)
              elif type == "loss":
                  self._loss_outputs.append(i)
              elif type == "variance":
                  self._variance_outputs.append(i)
              else:
                  self._other_outputs.append(i)
          if len(self._loss_outputs) == 0:
              self._loss_outputs = self._prediction_outputs
  
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 1 to create a new JaxModel object from inherited base classes: Model with example: obj = JaxModel(forward_fn, params, loss, output_types, batch_size, learning_rate, optimizer, grad_fn, update_fn, eval_fn, rng, log_frequency)