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: Create a generator that iterates batches for a dataset.
  Subclasses may override this method to customize how model inputs are
  generated from the data.
  
  Parameters
  ----------
  dataset: Dataset
      the data to iterate
  epochs: int
      the number of times to iterate over the full dataset
  mode: str
      allowed values are 'fit' (called during training), 'predict' (called
      during prediction), and 'uncertainty' (called during uncertainty
      prediction)
  deterministic: bool
      whether to iterate over the dataset in order, or randomly shuffle the
      data for each epoch
  pad_batches: bool
      whether to pad each batch up to this model's preferred batch size
  Returns
  -------
  a generator that iterates batches, each represented as a tuple of lists:
  ([inputs], [outputs], [weights])
docstr_tok: ["Create","a","generator","that","iterates","batches","for","a","dataset",".","Subclasses","may","override","this","method","to","customize","how","model","inputs","are","generated","from","the","data",".","Parameters","--","--","--","--","--","dataset",":","Dataset","the","data","to","iterate","epochs",":","int","the","number","of","times","to","iterate","over","the","full","dataset","mode",":","str","allowed","values","are","'fit","'","(","called","during","training",")",",","'predict","'","(","called","during","prediction",")",",","and","'uncertainty","'","(","called","during","uncertainty","prediction",")","deterministic",":","bool","whether","to","iterate","over","the","dataset","in","order",",","or","randomly","shuffle","the","data","for","each","epoch","pad_batches",":","bool","whether","to","pad","each","batch","up","to","this","model","'s","preferred","batch","size","Returns","--","--","--","-","a","generator","that","iterates","batches",",","each","represented","as","a","tuple","of","lists",":","(","[","inputs","]",",","[","outputs","]",",","[","weights","]",")"]
code: def default_generator(
      self,
      dataset: Dataset,
      epochs: int = 1,
      mode: str = "fit",
      deterministic: bool = True,
      pad_batches: bool = True,
  ) -> Iterable[Tuple[List, List, List]]:
      """Create a generator that iterates batches for a dataset.
      Subclasses may override this method to customize how model inputs are
      generated from the data.
  
      Parameters
      ----------
      dataset: Dataset
          the data to iterate
      epochs: int
          the number of times to iterate over the full dataset
      mode: str
          allowed values are 'fit' (called during training), 'predict' (called
          during prediction), and 'uncertainty' (called during uncertainty
          prediction)
      deterministic: bool
          whether to iterate over the dataset in order, or randomly shuffle the
          data for each epoch
      pad_batches: bool
          whether to pad each batch up to this model's preferred batch size
      Returns
      -------
      a generator that iterates batches, each represented as a tuple of lists:
      ([inputs], [outputs], [weights])
      """
  
      for epoch in range(epochs):
          for X_b, y_b, w_b, _ in dataset.iterbatches(
              batch_size=self.batch_size,
              deterministic=deterministic,
              pad_batches=pad_batches,
          ):
              yield ([X_b], [y_b], [w_b])
  
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 4 for calling the PINNModel obj's underlying member method code with example usage: obj.default_generator(dataset, epochs, mode, deterministic, pad_batches) without return types