class_docstr: Model provides an Estimator-like API similar to Keras.
docstr: Arguments:
      module: A `Module` instance.
      loss: A `elegy.Loss` or `Callable` instance representing the loss function of the network.
          You can define more loss terms by simply passing a possibly nested structure of
          lists and dictionaries of `elegy.Loss` or `Callable`s. Usually a plain list of losses is enough
          but using dictionaries will create namescopes for the names of the losses
          which might be useful e.g. to group things in tensorboard. Contrary to Keras convention,
          in Elegy there is no relation between the structure of `loss` with the structure
          of the labels and outputs of the network. Elegy's loss system is more flexible than
          the one provided by Keras, for more information on how to mimick Keras behavior checkout the
          [Losses and Metrics Guide](https://poets-ai.github.io/elegy/guides/losses-and-metrics)`.
      metrics: A `elegy.Metric` or `Callable` instance representing the loss function of the network.
          You can define more metrics terms by simply passing a possibly nested structure of
          lists and dictionaries of `elegy.Metric` or `Callable`s. Usually a plain list of metrics is enough
          but using dictionaries will create namescopes for the names of the metrics
          which might be useful e.g. to group things in tensorboard. Contrary to Keras convention,
          in Elegy there is no relation between the structure of `metrics` with the structure
          of the labels and outputs of the network. Elegy's metrics system is more flexible than
          the one provided by Keras, for more information on how to mimick Keras behavior checkout the
          [Losses and Metrics Guide](https://poets-ai.github.io/elegy/guides/losses-and-metrics)`.
      optimizer: A `optax` optimizer instance. Optix is a very flexible library for defining
          optimization pipelines with things like learning rate schedules, this means that
          there is no need for a `LearningRateScheduler` callback in Elegy.
      eager: Settable attribute indicating whether the model should run eagerly.
          Running eagerly means that your model will be run step by step, like Python code, instead of
          using Jax's `jit` to. Your model might run slower, but it should become easier for you to debug
          it by stepping into individual layer calls.
docstr_tok: ["Arguments",":","module",":","A","`","Module","`","instance",".","loss",":","A","`","elegy.Loss","`","or","`","Callable","`","instance","representing","the","loss","function","of","the","network",".","You","can","define","more","loss","terms","by","simply","passing","a","possibly","nested","structure","of","lists","and","dictionaries","of","`","elegy.Loss","`","or","`","Callable","`","s",".","Usually","a","plain","list","of","losses","is","enough","but","using","dictionaries","will","create","namescopes","for","the","names","of","the","losses","which","might","be","useful","e.g",".","to","group","things","in","tensorboard",".","Contrary","to","Keras","convention",",","in","Elegy","there","is","no","relation","between","the","structure","of","`","loss","`","with","the","structure","of","the","labels","and","outputs","of","the","network",".","Elegy","'s","loss","system","is","more","flexible","than","the","one","provided","by","Keras",",","for","more","information","on","how","to","mimick","Keras","behavior","checkout","the","[","Losses","and","Metrics","Guide","]","(","https",":","\/\/poets-ai.github.io\/elegy\/guides\/losses-and-metrics",")","`",".","metrics",":","A","`","elegy.Metric","`","or","`","Callable","`","instance","representing","the","loss","function","of","the","network",".","You","can","define","more","metrics","terms","by","simply","passing","a","possibly","nested","structure","of","lists","and","dictionaries","of","`","elegy.Metric","`","or","`","Callable","`","s",".","Usually","a","plain","list","of","metrics","is","enough","but","using","dictionaries","will","create","namescopes","for","the","names","of","the","metrics","which","might","be","useful","e.g",".","to","group","things","in","tensorboard",".","Contrary","to","Keras","convention",",","in","Elegy","there","is","no","relation","between","the","structure","of","`","metrics","`","with","the","structure","of","the","labels","and","outputs","of","the","network",".","Elegy","'s","metrics","system","is","more","flexible","than","the","one","provided","by","Keras",",","for","more","information","on","how","to","mimick","Keras","behavior","checkout","the","[","Losses","and","Metrics","Guide","]","(","https",":","\/\/poets-ai.github.io\/elegy\/guides\/losses-and-metrics",")","`",".","optimizer",":","A","`","optax","`","optimizer","instance",".","Optix","is","a","very","flexible","library","for","defining","optimization","pipelines","with","things","like","learning","rate","schedules",",","this","means","that","there","is","no","need","for","a","`","LearningRateScheduler","`","callback","in","Elegy",".","eager",":","Settable","attribute","indicating","whether","the","model","should","run","eagerly",".","Running","eagerly","means","that","your","model","will","be","run","step","by","step",",","like","Python","code",",","instead","of","using","Jax","'s","`","jit","`","to",".","Your","model","might","run","slower",",","but","it","should","become","easier","for","you","to","debug","it","by","stepping","into","individual","layer","calls","."]
code: def __init__(
      self,
      module: tp.Union[
          flax.linen.module.Module,
          TransformedWithState,
          tx.Module,
          None,
      ] = None,
      loss: tp.Any = None,
      metrics: tp.Any = None,
      optimizer: tp.Optional[
          tp.Union[tx.Optimizer, GradientTransformation]
      ] = None,
      seed: int = 42,
      eager: bool = False,
  ):
      """
      Arguments:
          module: A `Module` instance.
          loss: A `elegy.Loss` or `Callable` instance representing the loss function of the network.
              You can define more loss terms by simply passing a possibly nested structure of
              lists and dictionaries of `elegy.Loss` or `Callable`s. Usually a plain list of losses is enough
              but using dictionaries will create namescopes for the names of the losses
              which might be useful e.g. to group things in tensorboard. Contrary to Keras convention,
              in Elegy there is no relation between the structure of `loss` with the structure
              of the labels and outputs of the network. Elegy's loss system is more flexible than
              the one provided by Keras, for more information on how to mimick Keras behavior checkout the
              [Losses and Metrics Guide](https://poets-ai.github.io/elegy/guides/losses-and-metrics)`.
          metrics: A `elegy.Metric` or `Callable` instance representing the loss function of the network.
              You can define more metrics terms by simply passing a possibly nested structure of
              lists and dictionaries of `elegy.Metric` or `Callable`s. Usually a plain list of metrics is enough
              but using dictionaries will create namescopes for the names of the metrics
              which might be useful e.g. to group things in tensorboard. Contrary to Keras convention,
              in Elegy there is no relation between the structure of `metrics` with the structure
              of the labels and outputs of the network. Elegy's metrics system is more flexible than
              the one provided by Keras, for more information on how to mimick Keras behavior checkout the
              [Losses and Metrics Guide](https://poets-ai.github.io/elegy/guides/losses-and-metrics)`.
          optimizer: A `optax` optimizer instance. Optix is a very flexible library for defining
              optimization pipelines with things like learning rate schedules, this means that
              there is no need for a `LearningRateScheduler` callback in Elegy.
          eager: Settable attribute indicating whether the model should run eagerly.
              Running eagerly means that your model will be run step by step, like Python code, instead of
              using Jax's `jit` to. Your model might run slower, but it should become easier for you to debug
              it by stepping into individual layer calls.
      """
      super().__init__(seed=seed, eager=eager)
  
      if isinstance(module, flax.linen.module.Module):
          self.module = tx.FlaxModule(module)
  
      elif TransformedWithState is not None and isinstance(
          module, TransformedWithState
      ):
          self.module = HaikuModule(module)
      else:
          self.module = module
  
      self.optimizer = (
          tx.Optimizer(optimizer)
          if isinstance(optimizer, GradientTransformation)
          else optimizer
      )
      self.loss_and_logs = None
  
      self._losses_and_metrics = tx.Hashable((loss, metrics))
  
imports: import typing
  import flax
  import jax
  import jax.numpy
  import numpy
  import treex
  from optax import GradientTransformation
  from elegy import types, utils
  from elegy.model.model_base import ModelBase
  from elegy.model.model_core import GradStepOutput, PredStepOutput, TestStepOutput, TrainStepOutput
label_desc: Use image node_id 3 to create a new Model object from inherited base classes: tp.Generic, ModelBase with example: obj = Model(module, loss, metrics, optimizer, seed, eager)