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: Returns the loss value & metrics values for the model in test mode.
  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`.
      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`).
      verbose: 0, 1, or 2. Verbosity mode.
          0 = silent, 1 = progress bar, 2 = one line per epoch.
      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).
      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`.
      steps: Integer or `None`. Total number of steps (batches of samples)
          before declaring the evaluation round finished. Ignored with the
          default value of `None`. This
          argument is not supported with array inputs.
      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].
  
  Returns:
      A dictionary for mapping the losses and metrics names to the values obtained.
  Raises:
      ValueError: in case of invalid arguments.
docstr_tok: ["Returns","the","loss","value","&","metrics","values","for","the","model","in","test","mode",".","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","`",".","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","`",")",".","verbose",":","0",",","1",",","or","2",".","Verbosity","mode",".","0","=","silent",",","1","=","progress","bar",",","2","=","one","line","per","epoch",".","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",")",".","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","`",".","steps",":","Integer","or","`","None","`",".","Total","number","of","steps","(","batches","of","samples",")","before","declaring","the","evaluation","round","finished",".","Ignored","with","the","default","value","of","`","None","`",".","This","argument","is","not","supported","with","array","inputs",".","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","]",".","Returns",":","A","dictionary","for","mapping","the","losses","and","metrics","names","to","the","values","obtained",".","Raises",":","ValueError",":","in","case","of","invalid","arguments","."]
code: def evaluate(
      self,
      x: tp.Optional[tp.Any] = None,
      y: tp.Optional[tp.Any] = None,
      verbose: int = 1,
      batch_size: tp.Optional[int] = 32,
      sample_weight: tp.Optional[np.ndarray] = None,
      class_weight: tp.Optional[tp.Mapping[str, float]] = None,
      steps: tp.Optional[int] = None,
      callbacks: tp.Union[tp.List[Callback], CallbackList, None] = None,
      drop_remaining: bool = False,
  ) -> types.Logs:
      """Returns the loss value & metrics values for the model in test mode.
      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`.
          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`).
          verbose: 0, 1, or 2. Verbosity mode.
              0 = silent, 1 = progress bar, 2 = one line per epoch.
          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).
          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`.
          steps: Integer or `None`. Total number of steps (batches of samples)
              before declaring the evaluation round finished. Ignored with the
              default value of `None`. This
              argument is not supported with array inputs.
          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].
  
      Returns:
          A dictionary for mapping the losses and metrics names to the values obtained.
      Raises:
          ValueError: in case of invalid arguments.
      """
      if batch_size is not None:
          batch_size = self._distributed_strategy.lift_batch_size(
              batch_size
          )
  
      with _MODEL_CONTEXT.callbacks_context(
          self if _MODEL_CONTEXT.model is None else None
      ):
          if x is None:
              x = {}
  
          data_handler = data.DataHandler(
              x=x,
              y=y,
              sample_weight=sample_weight,
              batch_size=batch_size,
              steps_per_epoch=steps,
              initial_epoch=0,
              epochs=1,
              shuffle=False,
              training=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,
                  sigint_mode=SigIntMode.TEST,
                  model=self,
                  verbose=verbose,
                  epochs=1,
                  steps=data_handler.inferred_steps,
              )
  
          callbacks.on_test_begin()
  
          logs = {}
          for _, iterator in data_handler.enumerate_epochs():
              self.reset_metrics()
              with data_handler.catch_stop_iteration():
                  for step in data_handler.steps():
                      callbacks.on_test_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.test_on_batch(
                          inputs=inputs,
                          labels=labels,
                      )
                      tmp_logs.update({"size": data_handler.batch_size})
                      logs = tmp_logs
                      callbacks.on_test_batch_end(step, logs)
  
                      if self.stop_training:
                          break
              if self.stop_training:
                  break
  
          callbacks.on_test_end()
  
          return logs
  
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 6 for calling the ModelBase obj's underlying member method code with example usage: obj.evaluate(x, y, verbose, batch_size, sample_weight, class_weight, steps, callbacks, drop_remaining) and returns: logs