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: Generates output predictions for the input samples.
  Computation is done in batches.
  
  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 iterator types generator
          is given in the `Unpacking behavior for iterator-like inputs` section
          of `Model.fit`.
      batch_size: Integer or `None`.
          Number of samples per batch.
          If unspecified, `batch_size` will default to 32.
          Do not specify the `batch_size` if your data is in the
          form of generators (since they generate batches).
      verbose: Verbosity mode, 0 or 1.
      steps: Total number of steps (batches of samples)
          before declaring the prediction round finished.
          Ignored with the default value of `None`.
      callbacks: List of [elegy.callbacks.callback.Callback][] instances.
          List of callbacks to apply during training.
  
  See the discussion of `Unpacking behavior for iterator-like inputs` for
  [`Model.fit`][elegy.model.model.Model.fit].
  Note that Model.predict uses the same interpretation rules as
  `Model.fit` and `Model.evaluate`, so inputs must be unambiguous for all
  three methods.
  Returns:
      Numpy array(s) of predictions.
  Raises:
      ValueError: In case of mismatch between the provided
          input data and the model's expectations,
          or in case a stateful model receives a number of samples
          that is not a multiple of the batch size.
docstr_tok: ["Generates","output","predictions","for","the","input","samples",".","Computation","is","done","in","batches",".","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","iterator","types","generator","is","given","in","the","`","Unpacking","behavior","for","iterator-like","inputs","`","section","of","`","Model.fit","`",".","batch_size",":","Integer","or","`","None","`",".","Number","of","samples","per","batch",".","If","unspecified",",","`","batch_size","`","will","default","to","32",".","Do","not","specify","the","`","batch_size","`","if","your","data","is","in","the","form","of","generators","(","since","they","generate","batches",")",".","verbose",":","Verbosity","mode",",","0","or","1.","steps",":","Total","number","of","steps","(","batches","of","samples",")","before","declaring","the","prediction","round","finished",".","Ignored","with","the","default","value","of","`","None","`",".","callbacks",":","List","of","[","elegy.callbacks.callback.Callback","]","[","]","instances",".","List","of","callbacks","to","apply","during","training",".","See","the","discussion","of","`","Unpacking","behavior","for","iterator-like","inputs","`","for","[","`","Model.fit","`","]","[","elegy.model.model.Model.fit","]",".","Note","that","Model.predict","uses","the","same","interpretation","rules","as","`","Model.fit","`","and","`","Model.evaluate","`",",","so","inputs","must","be","unambiguous","for","all","three","methods",".","Returns",":","Numpy","array","(","s",")","of","predictions",".","Raises",":","ValueError",":","In","case","of","mismatch","between","the","provided","input","data","and","the","model","'s","expectations",",","or","in","case","a","stateful","model","receives","a","number","of","samples","that","is","not","a","multiple","of","the","batch","size","."]
code: def predict(
      self,
      x: tp.Optional[tp.Any] = None,
      verbose: int = 0,
      batch_size: tp.Optional[int] = 32,
      steps: tp.Optional[int] = None,
      callbacks: tp.Union[tp.List[Callback], CallbackList, None] = None,
      drop_remaining: bool = False,
  ) -> tp.Any:
      """Generates output predictions for the input samples.
      Computation is done in batches.
  
      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 iterator types generator
              is given in the `Unpacking behavior for iterator-like inputs` section
              of `Model.fit`.
          batch_size: Integer or `None`.
              Number of samples per batch.
              If unspecified, `batch_size` will default to 32.
              Do not specify the `batch_size` if your data is in the
              form of generators (since they generate batches).
          verbose: Verbosity mode, 0 or 1.
          steps: Total number of steps (batches of samples)
              before declaring the prediction round finished.
              Ignored with the default value of `None`.
          callbacks: List of [elegy.callbacks.callback.Callback][] instances.
              List of callbacks to apply during training.
  
      See the discussion of `Unpacking behavior for iterator-like inputs` for
      [`Model.fit`][elegy.model.model.Model.fit].
      Note that Model.predict uses the same interpretation rules as
      `Model.fit` and `Model.evaluate`, so inputs must be unambiguous for all
      three methods.
      Returns:
          Numpy array(s) of predictions.
      Raises:
          ValueError: In case of mismatch between the provided
              input data and the model's expectations,
              or in case a stateful model receives a number of samples
              that is not a multiple of the batch size.
      """
      if batch_size is not None:
          batch_size = self._distributed_strategy.lift_batch_size(
              batch_size
          )
  
      with _MODEL_CONTEXT.callbacks_context(self):
          if x is None:
              x = {}
  
          outputs = None
  
          data_handler = data.DataHandler(
              x=x,
              batch_size=batch_size,
              steps_per_epoch=steps,
              initial_epoch=0,
              epochs=1,
              shuffle=False,
          )
  
          # Container that configures and calls `tf.keras.Callback`s.
          if not isinstance(callbacks, CallbackList):
              callbacks = CallbackList(
                  callbacks,
                  add_history=True,
                  add_progbar=verbose != 0,
                  model=self,
                  verbose=verbose,
                  epochs=1,
                  steps=data_handler.inferred_steps,
              )
  
          callbacks.on_predict_begin()
  
          for _, iterator in data_handler.enumerate_epochs():
              self.reset_metrics()
              with data_handler.catch_stop_iteration():
                  for step in data_handler.steps():
                      callbacks.on_predict_batch_begin(step)
                      batch = next(iterator)
  
                      if (
                          drop_remaining
                          and not data_utils.has_batch_size(
                              batch, data_handler.batch_size
                          )
                      ):
                          continue
  
                      tmp_batch_outputs = self.predict_on_batch(
                          inputs=batch[0]
                      )
                      batch_outputs = tmp_batch_outputs
  
                      if outputs is None:
                          outputs = data.map_structure(
                              lambda batch_output: [batch_output],
                              batch_outputs,
                          )
                      else:
                          outputs = data.map_structure(
                              data.map_append,
                              outputs,
                              batch_outputs,
                          )
  
                      callbacks.on_predict_batch_end(
                          step,
                          {
                              "outputs": batch_outputs,
                              "size": data_handler.batch_size,
                          },
                      )
  
          callbacks.on_predict_end()
  
          all_outputs = data.map_structure(jnp.concatenate, outputs)
  
          return all_outputs
  
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 7 for calling the ModelBase obj's underlying member method code with example usage: obj.predict(x, verbose, batch_size, steps, callbacks, drop_remaining) and returns: all_outputs