class_docstr: `Model` is tasked with performing training, evaluation, and inference for a given
  `elegy.Module` or `haiku.Module`.
  
  To create a `Model` you first have to define its architecture in a `Module`:
  
  ```python
  >>> import elegy, jax
  >>> import jax.numpy as jnp
  
  >>> class MLP(elegy.Module):
  ...     def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
  ...         x = elegy.nn.Flatten()(x)
  ...         x = elegy.nn.Linear(5)(x)
  ...         x = jax.nn.relu(x)
  ...         x = elegy.nn.Linear(2)(x)
  ...         return x
  
  >>> mlp = MLP()
  >>> x = jnp.ones(shape=[10, 2])
  
  >>> y_pred, parameters, collections = mlp.init(rng=elegy.RNGSeq(42))(x)
  >>> y_pred.shape
  (10, 2)
  
  ```
  
  
  You can pass use `Module` with the Model API:
  ```python
  model = elegy.Model(
      module=MLP(),
      loss=[
          elegy.losses.Crossentropy(),
          elegy.regularizers.L2(l=1e-5),
      ],
      metrics=elegy.metrics.Accuracy(),
      optimizer=optax.rmsprop(1e-3),
  )
  ```
  
  Once the model is created, you can train the model with `model.fit()`, or use the model
  to do prediction with `model.predict()`.
  Checkout [Getting Started](https://poets-ai.github.io/elegy/getting-started) for
  additional details.
  
  Model supports defining + monitoring custom learning rate schedules by passing an instance of `elegy.Optimizer` instead of
  an `optax` object:
  
  ```python
  model = elegy.Model(
      module=MLP(n1=3, n2=1),
      loss=elegy.losses.Crossentropy(),
      metrics=elegy.metrics.Accuracy(),
      optimizer=elegy.Optimizer(
          optax.adam(1.0), # <---- important to set this to 1.0
          lr_schedule=lambda step, epoch: 1 / (epoch * 100 + step),
          steps_per_epoch=1000,
      ),
      eager=True,
  )
  
  history = model.fit(
      ...
  )
  
  assert "lr" in history.history
  ```
  Notice how we set the learning rate parameter of the `adam` optimizer to `1.0`, this is necessary if you want the logged `lr`
  be closer to the "actual" learning rate because we implement this feature by chaining an additional `optax.scale_by_schedule`
  at the end.
class_docstr_tok: ["`","Model","`","is","tasked","with","performing","training",",","evaluation",",","and","inference","for","a","given","`","elegy.Module","`","or","`","haiku.Module","`",".","To","create","a","`","Model","`","you","first","have","to","define","its","architecture","in","a","`","Module","`",":","``","`","python",">",">",">","import","elegy",",","jax",">",">",">","import","jax.numpy","as","jnp",">",">",">","class","MLP","(","elegy.Module",")",":","...","def","__call__","(","self",",","x",":","jnp.ndarray",")","-",">","jnp.ndarray",":","...","x","=","elegy.nn.Flatten","(",")","(","x",")","...","x","=","elegy.nn.Linear","(","5",")","(","x",")","...","x","=","jax.nn.relu","(","x",")","...","x","=","elegy.nn.Linear","(","2",")","(","x",")","...","return","x",">",">",">","mlp","=","MLP","(",")",">",">",">","x","=","jnp.ones","(","shape=","[","10",",","2","]",")",">",">",">","y_pred",",","parameters",",","collections","=","mlp.init","(","rng=elegy.RNGSeq","(","42",")",")","(","x",")",">",">",">","y_pred.shape","(","10",",","2",")","``","`","You","can","pass","use","`","Module","`","with","the","Model","API",":","``","`","python","model","=","elegy.Model","(","module=MLP","(",")",",","loss=","[","elegy.losses.Crossentropy","(",")",",","elegy.regularizers.L2","(","l=1e-5",")",",","]",",","metrics=elegy.metrics.Accuracy","(",")",",","optimizer=optax.rmsprop","(","1e-3",")",",",")","``","`","Once","the","model","is","created",",","you","can","train","the","model","with","`","model.fit","(",")","`",",","or","use","the","model","to","do","prediction","with","`","model.predict","(",")","`",".","Checkout","[","Getting","Started","]","(","https",":","\/\/poets-ai.github.io\/elegy\/getting-started",")","for","additional","details",".","Model","supports","defining","+","monitoring","custom","learning","rate","schedules","by","passing","an","instance","of","`","elegy.Optimizer","`","instead","of","an","`","optax","`","object",":","``","`","python","model","=","elegy.Model","(","module=MLP","(","n1=3",",","n2=1",")",",","loss=elegy.losses.Crossentropy","(",")",",","metrics=elegy.metrics.Accuracy","(",")",",","optimizer=elegy.Optimizer","(","optax.adam","(","1.0",")",",","#","<","--","--","important","to","set","this","to","1.0","lr_schedule=lambda","step",",","epoch",":","1","\/","(","epoch","*","100","+","step",")",",","steps_per_epoch=1000",",",")",",","eager=True",",",")","history","=","model.fit","(","...",")","assert","``","lr","''","in","history.history","``","`","Notice","how","we","set","the","learning","rate","parameter","of","the","`","adam","`","optimizer","to","`","1.0","`",",","this","is","necessary","if","you","want","the","logged","`","lr","`","be","closer","to","the","``","actual","''","learning","rate","because","we","implement","this","feature","by","chaining","an","additional","`","optax.scale_by_schedule","`","at","the","end","."]
docstr: Trains the model for a fixed number of epochs (iterations on a dataset).
  
  Arguments:
      x: Input data. It could be:
  
          - A Numpy or Jax array (or array-like), or a list of arrays
              (in case the model has multiple inputs).
          - A dict mapping input names to the corresponding arrays,
              if the model has named inputs.
          - A generator returning `(inputs,)`, `(inputs, targets)`
              or `(inputs, targets, sample_weights)`.
  
          A more detailed description of unpacking behavior for generator type
          is given below.
      y: Target data. Like the input data `x`,
          it could be either Numpy or Jax array(s).
          It should be consistent with `x`. If `x` is a generator,
          `y` should not be specified (since targets will be obtained from `x`).
      batch_size: Integer or `None`.
          Number of samples per gradient update.
          If unspecified, `batch_size` will default to 32.
          Do not specify the `batch_size` if your data is in the
          form of generator (since they generate batches).
      epochs: Integer. Number of epochs to train the model.
          An epoch is an iteration over the entire `x` and `y`
          data provided.
          Note that in conjunction with `initial_epoch`,
          `epochs` is to be understood as "final epoch".
          The model is not trained for a number of iterations
          given by `epochs`, but merely until the epoch
          of index `epochs` is reached.
      verbose: 0, 1, 2 or 3. Verbosity mode.
          0 = silent, 1 = progress bar, 2 = one line per epoch 3 = table.
          Note that the progress bar is not particularly useful when
          logged to a file, so verbose=2 is recommended when not running
          interactively (eg, in a production environment).
      callbacks: List of [elegy.callbacks.callback.Callback][] instances.
          List of callbacks to apply during training.
      validation_split: Float between 0 and 1.
          Fraction of the training data to be used as validation data.
          The model will set apart this fraction of the training data,
          will not train on it, and will evaluate
          the loss and any model metrics
          on this data at the end of each epoch.
          The validation data is selected from the last samples
          in the `x` and `y` data provided, before shuffling. This argument is
          not supported when `x` is a generator.
      validation_data: Data on which to evaluate
          the loss and any model metrics at the end of each epoch.
          The model will not be trained on this data.
          `validation_data` will override `validation_split`.
          `validation_data` could be:
  
          - tuple `(x_val, y_val)` of Numpy/Jax arrays, list of arrays or mappings
          - tuple `(x_val, y_val, val_sample_weights)` of Numpy/Jax arrays, list
          of arrays or mappings
          - generator
  
          For the first two cases, `batch_size` must be provided.
          For the last case, `validation_steps` should be provided, and should
          follow the same convention for yielding data as `x`.
          Note that `validation_data` does not support all the data types that
          are supported in `x`, eg, dict.
      shuffle: Boolean (whether to shuffle the training data
          before each epoch). This argument is ignored
          when `x` is a generator. Has no effect when `steps_per_epoch` is not `None`.
      class_weight: Optional dictionary mapping class indices (integers)
          to a weight (float) value, used for weighting the loss function
          (during training only).
          This can be useful to tell the model to
          "pay more attention" to samples from
          an under-represented class.
      sample_weight: Optional Numpy/Jax array of weights for
          the training samples, used for weighting the loss function
          (during training only). You can either pass a flat (1D)
          Numpy array with the same length as the input samples
          (1:1 mapping between weights and samples). This argument is not
          supported when `x` is generator, instead provide the sample_weights
          as the third element of `x`.
      initial_epoch: Integer.
          Epoch at which to start training
          (useful for resuming a previous training run).
      steps_per_epoch: Integer or `None`.
          Total number of steps (batches of samples)
          before declaring one epoch finished and starting the
          next epoch. When training with input arrays such as
          jax data arrays, the default `None` is equal to
          the number of samples in your dataset divided by
          the batch size, or 1 if that cannot be determined.
          When passing a generator, you must specify the
          `steps_per_epoch` argument. This argument is not supported with
          array inputs.
      validation_steps: Only relevant if `validation_data` is provided and
          is a generator. Total number of steps (batches of
          samples) to draw before stopping when performing validation
          at the end of every epoch. If 'validation_steps' is None, validation
          will run until the `validation_data` dataset is exhausted. In the
          case of an infinitely repeated dataset, it will run into an
          infinite loop. If 'validation_steps' is specified and only part of
          the dataset will be consumed, the evaluation will start from the
          beginning of the dataset at each epoch. This ensures that the same
          validation samples are used every time.
      validation_batch_size: Integer or `None`.
          Number of samples per validation batch.
          If unspecified, will default to `batch_size`.
          Do not specify the `validation_batch_size` if your data is in the
          form of generators (since they generate batches).
      validation_freq: Only relevant if validation data is provided. Integer
          or `collections_abc.Container` instance (e.g. list, tuple, etc.).
          If an integer, specifies how many training epochs to run before a
          new validation run is performed, e.g. `validation_freq=2` runs
          validation every 2 epochs. If a Container, specifies the epochs on
          which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
          validation at the end of the 1st, 2nd, and 10th epochs.
  
  Unpacking behavior for iterator-like inputs:
  
  A common pattern is to pass a generator, which will in fact
  yield not only features (x) but optionally targets (y) and sample weights.
  Elegy requires that the output of such iterator-likes be unambiguous. The
  iterator should return a tuple of length 1, 2, or 3, where the optional
  second and third elements will be used for y and sample_weight
  respectively. Any other type provided will be wrapped in a length one
  tuple, effectively treating everything as 'x'. When yielding dicts, they
  should still adhere to the top-level tuple structure.
  e.g. `({"x0": x0, "x1": x1}, y)`. Elegy will not attempt to separate
  features, targets, and weights from the keys of a single dict.
  
  A notable unsupported data type is the namedtuple. The reason is that
  it behaves like both an ordered datatype (tuple) and a mapping
  datatype (dict). So given a namedtuple of the form:
      `namedtuple("example_tuple", ["y", "x"])`
  it is ambiguous whether to reverse the order of the elements when
  interpreting the value. Even worse is a tuple of the form:
      `namedtuple("other_tuple", ["x", "y", "z"])`
  where it is unclear if the tuple was intended to be unpacked into x, y,
  and sample_weight or passed through as a single element to `x`. As a
  result the data processing code will simply raise a ValueError if it
  encounters a namedtuple. (Along with instructions to remedy the issue.)
  
  Returns:
      A `History` object. Its `History.history` attribute is
      a record of training loss values and metrics values
      at successive epochs, as well as validation loss values
      and validation metrics values (if applicable).
  Raises:
      ValueError: In case of mismatch between the provided input data
          and what the model expects.
docstr_tok: ["Trains","the","model","for","a","fixed","number","of","epochs","(","iterations","on","a","dataset",")",".","Arguments",":","x",":","Input","data",".","It","could","be",":","-","A","Numpy","or","Jax","array","(","or","array-like",")",",","or","a","list","of","arrays","(","in","case","the","model","has","multiple","inputs",")",".","-","A","dict","mapping","input","names","to","the","corresponding","arrays",",","if","the","model","has","named","inputs",".","-","A","generator","returning","`","(","inputs",",",")","`",",","`","(","inputs",",","targets",")","`","or","`","(","inputs",",","targets",",","sample_weights",")","`",".","A","more","detailed","description","of","unpacking","behavior","for","generator","type","is","given","below",".","y",":","Target","data",".","Like","the","input","data","`","x","`",",","it","could","be","either","Numpy","or","Jax","array","(","s",")",".","It","should","be","consistent","with","`","x","`",".","If","`","x","`","is","a","generator",",","`","y","`","should","not","be","specified","(","since","targets","will","be","obtained","from","`","x","`",")",".","batch_size",":","Integer","or","`","None","`",".","Number","of","samples","per","gradient","update",".","If","unspecified",",","`","batch_size","`","will","default","to","32",".","Do","not","specify","the","`","batch_size","`","if","your","data","is","in","the","form","of","generator","(","since","they","generate","batches",")",".","epochs",":","Integer",".","Number","of","epochs","to","train","the","model",".","An","epoch","is","an","iteration","over","the","entire","`","x","`","and","`","y","`","data","provided",".","Note","that","in","conjunction","with","`","initial_epoch","`",",","`","epochs","`","is","to","be","understood","as","``","final","epoch","''",".","The","model","is","not","trained","for","a","number","of","iterations","given","by","`","epochs","`",",","but","merely","until","the","epoch","of","index","`","epochs","`","is","reached",".","verbose",":","0",",","1",",","2","or","3",".","Verbosity","mode",".","0","=","silent",",","1","=","progress","bar",",","2","=","one","line","per","epoch","3","=","table",".","Note","that","the","progress","bar","is","not","particularly","useful","when","logged","to","a","file",",","so","verbose=2","is","recommended","when","not","running","interactively","(","eg",",","in","a","production","environment",")",".","callbacks",":","List","of","[","elegy.callbacks.callback.Callback","]","[","]","instances",".","List","of","callbacks","to","apply","during","training",".","validation_split",":","Float","between","0","and","1",".","Fraction","of","the","training","data","to","be","used","as","validation","data",".","The","model","will","set","apart","this","fraction","of","the","training","data",",","will","not","train","on","it",",","and","will","evaluate","the","loss","and","any","model","metrics","on","this","data","at","the","end","of","each","epoch",".","The","validation","data","is","selected","from","the","last","samples","in","the","`","x","`","and","`","y","`","data","provided",",","before","shuffling",".","This","argument","is","not","supported","when","`","x","`","is","a","generator",".","validation_data",":","Data","on","which","to","evaluate","the","loss","and","any","model","metrics","at","the","end","of","each","epoch",".","The","model","will","not","be","trained","on","this","data",".","`","validation_data","`","will","override","`","validation_split","`",".","`","validation_data","`","could","be",":","-","tuple","`","(","x_val",",","y_val",")","`","of","Numpy\/Jax","arrays",",","list","of","arrays","or","mappings","-","tuple","`","(","x_val",",","y_val",",","val_sample_weights",")","`","of","Numpy\/Jax","arrays",",","list","of","arrays","or","mappings","-","generator","For","the","first","two","cases",",","`","batch_size","`","must","be","provided",".","For","the","last","case",",","`","validation_steps","`","should","be","provided",",","and","should","follow","the","same","convention","for","yielding","data","as","`","x","`",".","Note","that","`","validation_data","`","does","not","support","all","the","data","types","that","are","supported","in","`","x","`",",","eg",",","dict",".","shuffle",":","Boolean","(","whether","to","shuffle","the","training","data","before","each","epoch",")",".","This","argument","is","ignored","when","`","x","`","is","a","generator",".","Has","no","effect","when","`","steps_per_epoch","`","is","not","`","None","`",".","class_weight",":","Optional","dictionary","mapping","class","indices","(","integers",")","to","a","weight","(","float",")","value",",","used","for","weighting","the","loss","function","(","during","training","only",")",".","This","can","be","useful","to","tell","the","model","to","``","pay","more","attention","''","to","samples","from","an","under-represented","class",".","sample_weight",":","Optional","Numpy\/Jax","array","of","weights","for","the","training","samples",",","used","for","weighting","the","loss","function","(","during","training","only",")",".","You","can","either","pass","a","flat","(","1D",")","Numpy","array","with","the","same","length","as","the","input","samples","(","1:1","mapping","between","weights","and","samples",")",".","This","argument","is","not","supported","when","`","x","`","is","generator",",","instead","provide","the","sample_weights","as","the","third","element","of","`","x","`",".","initial_epoch",":","Integer",".","Epoch","at","which","to","start","training","(","useful","for","resuming","a","previous","training","run",")",".","steps_per_epoch",":","Integer","or","`","None","`",".","Total","number","of","steps","(","batches","of","samples",")","before","declaring","one","epoch","finished","and","starting","the","next","epoch",".","When","training","with","input","arrays","such","as","jax","data","arrays",",","the","default","`","None","`","is","equal","to","the","number","of","samples","in","your","dataset","divided","by","the","batch","size",",","or","1","if","that","can","not","be","determined",".","When","passing","a","generator",",","you","must","specify","the","`","steps_per_epoch","`","argument",".","This","argument","is","not","supported","with","array","inputs",".","validation_steps",":","Only","relevant","if","`","validation_data","`","is","provided","and","is","a","generator",".","Total","number","of","steps","(","batches","of","samples",")","to","draw","before","stopping","when","performing","validation","at","the","end","of","every","epoch",".","If","'validation_steps","'","is","None",",","validation","will","run","until","the","`","validation_data","`","dataset","is","exhausted",".","In","the","case","of","an","infinitely","repeated","dataset",",","it","will","run","into","an","infinite","loop",".","If","'validation_steps","'","is","specified","and","only","part","of","the","dataset","will","be","consumed",",","the","evaluation","will","start","from","the","beginning","of","the","dataset","at","each","epoch",".","This","ensures","that","the","same","validation","samples","are","used","every","time",".","validation_batch_size",":","Integer","or","`","None","`",".","Number","of","samples","per","validation","batch",".","If","unspecified",",","will","default","to","`","batch_size","`",".","Do","not","specify","the","`","validation_batch_size","`","if","your","data","is","in","the","form","of","generators","(","since","they","generate","batches",")",".","validation_freq",":","Only","relevant","if","validation","data","is","provided",".","Integer","or","`","collections_abc.Container","`","instance","(","e.g",".","list",",","tuple",",","etc.",")",".","If","an","integer",",","specifies","how","many","training","epochs","to","run","before","a","new","validation","run","is","performed",",","e.g",".","`","validation_freq=2","`","runs","validation","every","2","epochs",".","If","a","Container",",","specifies","the","epochs","on","which","to","run","validation",",","e.g",".","`","validation_freq=","[","1",",","2",",","10","]","`","runs","validation","at","the","end","of","the","1st",",","2nd",",","and","10th","epochs",".","Unpacking","behavior","for","iterator-like","inputs",":","A","common","pattern","is","to","pass","a","generator",",","which","will","in","fact","yield","not","only","features","(","x",")","but","optionally","targets","(","y",")","and","sample","weights",".","Elegy","requires","that","the","output","of","such","iterator-likes","be","unambiguous",".","The","iterator","should","return","a","tuple","of","length","1",",","2",",","or","3",",","where","the","optional","second","and","third","elements","will","be","used","for","y","and","sample_weight","respectively",".","Any","other","type","provided","will","be","wrapped","in","a","length","one","tuple",",","effectively","treating","everything","as","'","x","'",".","When","yielding","dicts",",","they","should","still","adhere","to","the","top-level","tuple","structure",".","e.g",".","`","(","{","``","x0","''",":","x0",",","``","x1","''",":","x1","}",",","y",")","`",".","Elegy","will","not","attempt","to","separate","features",",","targets",",","and","weights","from","the","keys","of","a","single","dict",".","A","notable","unsupported","data","type","is","the","namedtuple",".","The","reason","is","that","it","behaves","like","both","an","ordered","datatype","(","tuple",")","and","a","mapping","datatype","(","dict",")",".","So","given","a","namedtuple","of","the","form",":","`","namedtuple","(","``","example_tuple","''",",","[","``","y","''",",","``","x","''","]",")","`","it","is","ambiguous","whether","to","reverse","the","order","of","the","elements","when","interpreting","the","value",".","Even","worse","is","a","tuple","of","the","form",":","`","namedtuple","(","``","other_tuple","''",",","[","``","x","''",",","``","y","''",",","``","z","''","]",")","`","where","it","is","unclear","if","the","tuple","was","intended","to","be","unpacked","into","x",",","y",",","and","sample_weight","or","passed","through","as","a","single","element","to","`","x","`",".","As","a","result","the","data","processing","code","will","simply","raise","a","ValueError","if","it","encounters","a","namedtuple",".","(","Along","with","instructions","to","remedy","the","issue",".",")","Returns",":","A","`","History","`","object",".","Its","`","History.history","`","attribute","is","a","record","of","training","loss","values","and","metrics","values","at","successive","epochs",",","as","well","as","validation","loss","values","and","validation","metrics","values","(","if","applicable",")",".","Raises",":","ValueError",":","In","case","of","mismatch","between","the","provided","input","data","and","what","the","model","expects","."]
code: def fit(
      self,
      inputs: tp.Optional[tp.Any] = None,
      labels: tp.Optional[tp.Any] = None,
      batch_size: tp.Optional[int] = 32,
      epochs: int = 1,
      verbose: int = 1,
      callbacks: tp.Union[tp.List[Callback], CallbackList, None] = None,
      validation_split: float = 0.0,
      validation_data: tp.Union[tp.Tuple, tp.Iterable, None] = None,
      shuffle: bool = True,
      initial_epoch: int = 0,
      steps_per_epoch: tp.Optional[int] = None,
      validation_steps: tp.Optional[int] = None,
      validation_batch_size: tp.Optional[int] = None,
      validation_freq: int = 1,
      drop_remaining: bool = True,
  ) -> History:
      """
      Trains the model for a fixed number of epochs (iterations on a dataset).
  
      Arguments:
          x: Input data. It could be:
  
              - A Numpy or Jax array (or array-like), or a list of arrays
                  (in case the model has multiple inputs).
              - A dict mapping input names to the corresponding arrays,
                  if the model has named inputs.
              - A generator returning `(inputs,)`, `(inputs, targets)`
                  or `(inputs, targets, sample_weights)`.
  
              A more detailed description of unpacking behavior for generator type
              is given below.
          y: Target data. Like the input data `x`,
              it could be either Numpy or Jax array(s).
              It should be consistent with `x`. If `x` is a generator,
              `y` should not be specified (since targets will be obtained from `x`).
          batch_size: Integer or `None`.
              Number of samples per gradient update.
              If unspecified, `batch_size` will default to 32.
              Do not specify the `batch_size` if your data is in the
              form of generator (since they generate batches).
          epochs: Integer. Number of epochs to train the model.
              An epoch is an iteration over the entire `x` and `y`
              data provided.
              Note that in conjunction with `initial_epoch`,
              `epochs` is to be understood as "final epoch".
              The model is not trained for a number of iterations
              given by `epochs`, but merely until the epoch
              of index `epochs` is reached.
          verbose: 0, 1, 2 or 3. Verbosity mode.
              0 = silent, 1 = progress bar, 2 = one line per epoch 3 = table.
              Note that the progress bar is not particularly useful when
              logged to a file, so verbose=2 is recommended when not running
              interactively (eg, in a production environment).
          callbacks: List of [elegy.callbacks.callback.Callback][] instances.
              List of callbacks to apply during training.
          validation_split: Float between 0 and 1.
              Fraction of the training data to be used as validation data.
              The model will set apart this fraction of the training data,
              will not train on it, and will evaluate
              the loss and any model metrics
              on this data at the end of each epoch.
              The validation data is selected from the last samples
              in the `x` and `y` data provided, before shuffling. This argument is
              not supported when `x` is a generator.
          validation_data: Data on which to evaluate
              the loss and any model metrics at the end of each epoch.
              The model will not be trained on this data.
              `validation_data` will override `validation_split`.
              `validation_data` could be:
  
              - tuple `(x_val, y_val)` of Numpy/Jax arrays, list of arrays or mappings
              - tuple `(x_val, y_val, val_sample_weights)` of Numpy/Jax arrays, list
              of arrays or mappings
              - generator
  
              For the first two cases, `batch_size` must be provided.
              For the last case, `validation_steps` should be provided, and should
              follow the same convention for yielding data as `x`.
              Note that `validation_data` does not support all the data types that
              are supported in `x`, eg, dict.
          shuffle: Boolean (whether to shuffle the training data
              before each epoch). This argument is ignored
              when `x` is a generator. Has no effect when `steps_per_epoch` is not `None`.
          class_weight: Optional dictionary mapping class indices (integers)
              to a weight (float) value, used for weighting the loss function
              (during training only).
              This can be useful to tell the model to
              "pay more attention" to samples from
              an under-represented class.
          sample_weight: Optional Numpy/Jax array of weights for
              the training samples, used for weighting the loss function
              (during training only). You can either pass a flat (1D)
              Numpy array with the same length as the input samples
              (1:1 mapping between weights and samples). This argument is not
              supported when `x` is generator, instead provide the sample_weights
              as the third element of `x`.
          initial_epoch: Integer.
              Epoch at which to start training
              (useful for resuming a previous training run).
          steps_per_epoch: Integer or `None`.
              Total number of steps (batches of samples)
              before declaring one epoch finished and starting the
              next epoch. When training with input arrays such as
              jax data arrays, the default `None` is equal to
              the number of samples in your dataset divided by
              the batch size, or 1 if that cannot be determined.
              When passing a generator, you must specify the
              `steps_per_epoch` argument. This argument is not supported with
              array inputs.
          validation_steps: Only relevant if `validation_data` is provided and
              is a generator. Total number of steps (batches of
              samples) to draw before stopping when performing validation
              at the end of every epoch. If 'validation_steps' is None, validation
              will run until the `validation_data` dataset is exhausted. In the
              case of an infinitely repeated dataset, it will run into an
              infinite loop. If 'validation_steps' is specified and only part of
              the dataset will be consumed, the evaluation will start from the
              beginning of the dataset at each epoch. This ensures that the same
              validation samples are used every time.
          validation_batch_size: Integer or `None`.
              Number of samples per validation batch.
              If unspecified, will default to `batch_size`.
              Do not specify the `validation_batch_size` if your data is in the
              form of generators (since they generate batches).
          validation_freq: Only relevant if validation data is provided. Integer
              or `collections_abc.Container` instance (e.g. list, tuple, etc.).
              If an integer, specifies how many training epochs to run before a
              new validation run is performed, e.g. `validation_freq=2` runs
              validation every 2 epochs. If a Container, specifies the epochs on
              which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
              validation at the end of the 1st, 2nd, and 10th epochs.
  
      Unpacking behavior for iterator-like inputs:
  
      A common pattern is to pass a generator, which will in fact
      yield not only features (x) but optionally targets (y) and sample weights.
      Elegy requires that the output of such iterator-likes be unambiguous. The
      iterator should return a tuple of length 1, 2, or 3, where the optional
      second and third elements will be used for y and sample_weight
      respectively. Any other type provided will be wrapped in a length one
      tuple, effectively treating everything as 'x'. When yielding dicts, they
      should still adhere to the top-level tuple structure.
      e.g. `({"x0": x0, "x1": x1}, y)`. Elegy will not attempt to separate
      features, targets, and weights from the keys of a single dict.
  
      A notable unsupported data type is the namedtuple. The reason is that
      it behaves like both an ordered datatype (tuple) and a mapping
      datatype (dict). So given a namedtuple of the form:
          `namedtuple("example_tuple", ["y", "x"])`
      it is ambiguous whether to reverse the order of the elements when
      interpreting the value. Even worse is a tuple of the form:
          `namedtuple("other_tuple", ["x", "y", "z"])`
      where it is unclear if the tuple was intended to be unpacked into x, y,
      and sample_weight or passed through as a single element to `x`. As a
      result the data processing code will simply raise a ValueError if it
      encounters a namedtuple. (Along with instructions to remedy the issue.)
  
      Returns:
          A `History` object. Its `History.history` attribute is
          a record of training loss values and metrics values
          at successive epochs, as well as validation loss values
          and validation metrics values (if applicable).
      Raises:
          ValueError: In case of mismatch between the provided input data
              and what the model expects.
      """
  
      if batch_size is not None:
          batch_size = self._distributed_strategy.lift_batch_size(
              batch_size
          )
  
      with _MODEL_CONTEXT.callbacks_context(self):
          if inputs is None:
              inputs = dict()
  
          if validation_split:
              # Create the validation data using the training data. Only supported for
              # `Jax Numpy` and `NumPy` input.
              (
                  inputs,
                  labels,
              ), validation_data = data.train_validation_split(
                  (inputs, labels),
                  validation_split=validation_split,
                  shuffle=False,
              )
  
          self.stop_training = False
          data_handler = data.DataHandler(
              x=inputs,
              y=labels,
              # sample_weight=sample_weight,
              batch_size=batch_size,
              steps_per_epoch=steps_per_epoch,
              initial_epoch=initial_epoch,
              epochs=epochs,
              shuffle=shuffle,
              # class_weight=class_weight,
          )
          # Container that configures and calls `tf.keras.Callback`s.
          if not isinstance(callbacks, CallbackList):
              callbacks = CallbackList(
                  callbacks,
                  add_history=True,
                  add_progbar=verbose != 0,
                  sigint_mode=SigIntMode.TRAIN,
                  model=self,
                  verbose=verbose,
                  epochs=epochs,
                  steps=data_handler.inferred_steps,
              )
  
          callbacks.on_train_begin()
          # data_handler._initial_epoch = (  # pylint: disable=protected-access
          #     self._maybe_load_initial_epoch_from_ckpt(initial_epoch))
          epoch_logs = {}
  
          for epoch, iterator in data_handler.enumerate_epochs():
              self.reset_metrics()
              callbacks.on_epoch_begin(epoch)
              logs = {}
              with data_handler.catch_stop_iteration():
                  for step in data_handler.steps():
                      # prev_model = self.copy()
  
                      callbacks.on_train_batch_begin(step)
                      batch = next(iterator)
                      inputs, labels, _ = data.unpack_x_y_sample_weight(
                          batch
                      )
  
                      if (
                          drop_remaining
                          and not data_utils.has_batch_size(
                              batch, data_handler.batch_size
                          )
                      ):
                          continue
  
                      tmp_logs = self.train_on_batch(
                          inputs=inputs,
                          labels=labels,
                      )
                      tmp_logs.update({"size": data_handler.batch_size})
  
                      logs = tmp_logs
                      callbacks.on_train_batch_end(step, logs)
  
                      if self.stop_training:
                          break
  
                      # utils._walk_treedef(
                      #     jax.tree_flatten(prev_model)[1],
                      #     jax.tree_flatten(self)[1],
                      # )
  
              epoch_logs = copy(logs)
              epoch_logs.update({"size": data_handler.batch_size})
  
              # Run validation.
              if (
                  validation_data
                  and self._should_eval(epoch, validation_freq)
                  and not self.stop_training
              ):
                  (
                      val_inputs,
                      val_lables,
                      _,
                  ) = data.unpack_x_y_sample_weight(validation_data)
                  # val_inputs, val_lables = validation_data
                  try:
                      val_logs = self.evaluate(
                          x=val_inputs,
                          y=val_lables,
                          batch_size=validation_batch_size
                          or batch_size,
                          steps=validation_steps,
                          callbacks=callbacks,
                          # return_dict=True,
                          drop_remaining=drop_remaining,
                      )
  
                      val_logs = {
                          "val_" + name: val
                          for name, val in val_logs.items()
                      }
                      epoch_logs.update(val_logs)
                  except (
                      types.MissingMethod,
                      types.MissingModule,
                  ) as e:
                      pass
  
              callbacks.on_epoch_end(epoch, epoch_logs)
  
              if self.stop_training:
                  break
  
          callbacks.on_train_end(epoch_logs)
  
          assert self.history is not None
  
          return self.history
  
imports: import pickle
  import threading
  import typing
  from contextlib import contextmanager
  from copy import copy
  from dataclasses import dataclass
  from pathlib import Path
  import jax
  import jax.numpy
  import numpy
  import treex
  from elegy import data, types, utils
  from elegy.callbacks import Callback, CallbackList, History, history
  from elegy.callbacks.sigint import SigIntMode
  from elegy.data import utils
  from elegy.model.model_core import ModelCore, PredStepOutput, TestStepOutput
label_desc: Use image node_id 5 for calling the ModelBase obj's underlying member method code with example usage: obj.fit(inputs, labels, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, drop_remaining) and returns: self