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","."]
code: def stop_training(self) -> bool:
      if _MODEL_CONTEXT.model is not self:
          raise ValueError(
              f"Trying to get `.history` outside `callbacks_context`"
          )
  
      return _MODEL_CONTEXT.stop_training
  
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 3 for calling the ModelBase obj's underlying member method code with example usage: obj.stop_training() and returns: _MODEL_CONTEXT