class_docstr: This is class is derived from the JaxModel class and methods are also very similar to JaxModel,
  but it has the option of passing multiple arguments(Done using *args) suitable for PINNs model.
  Ex - Approximating f(x, y, z, t) satisfying a Linear differential equation.
  
  This model is recommended for linear partial differential equations but if you can accurately write
  the gradient function in Jax depending on your use case, then it will work as well.
  
  This class requires two functions apart from the usual function definition and weights
  
  [1] **grad_fn** : Each PINNs have a different strategy for calculating its final losses. This
  function tells the PINNModel how to go about computing the derivatives for backpropagation.
  It should follow this format:
  
  >>>
  >> def gradient_fn(forward_fn, loss_outputs, initial_data):
  >>
  >>  def model_loss(params, target, weights, rng, ...):
  >>
  >>    # write code using the arguments.
  >>    # ... indicates the variable number of positional arguments.
  >>    return
  >>
  >>  return model_loss
  
  "..." can be replaced with various arguments like (x, y, z, y) but should match with eval_fn
  
  [2] **eval_fn**: Function for defining how the model needs to compute during inference.
  It should follow this format
  
  >>>
  >> def create_eval_fn(forward_fn, params):
  >>  def eval_model(..., rng=None):
  >>    # write code here using arguments
  >>
  >>    return
  >>  return eval_model
  
  "..." can be replaced with various arguments like (x, y, z, y) but should match with grad_fn
  
  [3] boundary_data:
  For a detailed example, check out - deepchem/models/jax_models/tests/test_pinn.py where we have
  solved f'(x) = -sin(x)
  
  References
  ----------
  .. [1] Raissi et. al. "Physics-informed neural networks: A deep learning framework for solving
      forward and inverse problems involving nonlinear partial differential equations" Journal of
      Computational Physics https://doi.org/10.1016/j.jcp.2018.10.045
  
  .. [2] Raissi et. al. "Physics Informed Deep Learning (Part I): Data-driven
      Solutions of Nonlinear Partial Differential Equations" arXiv preprint arXiv:1711.10561
  
  Notes
  -----
  This class requires Jax, Haiku and Optax to be installed.
class_docstr_tok: ["This","is","class","is","derived","from","the","JaxModel","class","and","methods","are","also","very","similar","to","JaxModel",",","but","it","has","the","option","of","passing","multiple","arguments","(","Done","using","*","args",")","suitable","for","PINNs","model",".","Ex","-","Approximating","f","(","x",",","y",",","z",",","t",")","satisfying","a","Linear","differential","equation",".","This","model","is","recommended","for","linear","partial","differential","equations","but","if","you","can","accurately","write","the","gradient","function","in","Jax","depending","on","your","use","case",",","then","it","will","work","as","well",".","This","class","requires","two","functions","apart","from","the","usual","function","definition","and","weights","[","1","]","*","*","grad_fn","*","*",":","Each","PINNs","have","a","different","strategy","for","calculating","its","final","losses",".","This","function","tells","the","PINNModel","how","to","go","about","computing","the","derivatives","for","backpropagation",".","It","should","follow","this","format",":",">",">",">",">",">","def","gradient_fn","(","forward_fn",",","loss_outputs",",","initial_data",")",":",">",">",">",">","def","model_loss","(","params",",","target",",","weights",",","rng",",","...",")",":",">",">",">",">","#","write","code","using","the","arguments",".",">",">","#","...","indicates","the","variable","number","of","positional","arguments",".",">",">","return",">",">",">",">","return","model_loss","''","...","''","can","be","replaced","with","various","arguments","like","(","x",",","y",",","z",",","y",")","but","should","match","with","eval_fn","[","2","]","*","*","eval_fn","*","*",":","Function","for","defining","how","the","model","needs","to","compute","during","inference",".","It","should","follow","this","format",">",">",">",">",">","def","create_eval_fn","(","forward_fn",",","params",")",":",">",">","def","eval_model","(","...",",","rng=None",")",":",">",">","#","write","code","here","using","arguments",">",">",">",">","return",">",">","return","eval_model","''","...","''","can","be","replaced","with","various","arguments","like","(","x",",","y",",","z",",","y",")","but","should","match","with","grad_fn","[","3","]","boundary_data",":","For","a","detailed","example",",","check","out","-","deepchem\/models\/jax_models\/tests\/test_pinn.py","where","we","have","solved","f","'","(","x",")","=","-sin","(","x",")","References","--","--","--","--","--","..","[","1","]","Raissi","et",".","al",".","``","Physics-informed","neural","networks",":","A","deep","learning","framework","for","solving","forward","and","inverse","problems","involving","nonlinear","partial","differential","equations","''","Journal","of","Computational","Physics","https",":","\/\/doi.org\/10.1016\/j.jcp.2018.10.045","..","[","2","]","Raissi","et",".","al",".","``","Physics","Informed","Deep","Learning","(","Part","I",")",":","Data-driven","Solutions","of","Nonlinear","Partial","Differential","Equations","''","arXiv","preprint","arXiv:1711.10561","Notes","--","--","-","This","class","requires","Jax",",","Haiku","and","Optax","to","be","installed","."]
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.
docstr_tok: ["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
  from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
  from collections.abc import Sequence
  from deepchem.data import Dataset
  from deepchem.models.losses import Loss
  from deepchem.models.optimizers import Optimizer
  from deepchem.utils.typing import LossFn, OneOrMany
  from deepchem.trans.transformers import Transformer, undo_transforms
  import jax
  import jax.numpy
  import haiku
  import optax
  from deepchem.models.jax_models.jax_model import JaxModel
  from deepchem.models.jax_models.jax_model import create_default_gradient_fn, create_default_eval_fn
  import logging
  import warnings
label_desc: Use image node_id 5 for calling the PINNModel obj's underlying member method code with example usage: obj._predict(generator, transformers, uncertainty, other_output_types) and returns: zip, final_results, final_results