flax_extra.console module

Logging to cosole.

flax_extra.console.log(message: str, stdout: bool) None[source]

Writes a message to stdout.

Parameters
  • message – an arbitrary text.

  • stdout – whether to write or not to stdout.

flax_extra.console.log_step(step: int, message: str, stdout: bool = True) None[source]

Writes a step-related message to stdout.

Parameters
  • step – the step number.

  • message – an arbitrary text.

  • stdout – whether to write or not to stdout.