Welcome to Jaxer’s documentation!

Jax and Flax Time Series Prediction Transformer. The initial goal of this repo was to learn Jax and Flax by implementing a deep learning model. In this case, a transformer for time series prediction. I have decided to predict cryptocurrency due to the availability of the data. However, this could be used for any other time series data, such as market stocks or energy demand.

Note

This doc has been generated not only to show my work but also to document a real example on how to use Jax and Flax to build a sota deep learning model from scratch. Hope you find this library useful. Feel free to contribute to the project and make it better. You are welcome to open an issue or a pull request!

In order to solve cryptocurrency prediction, several approaches are proposed:

• Mean prediction: Predict the next value of the time series (mean).

• Distribution prediction: Predict the distribution over the next step (mean and log std) to cover the uncertainty.

• Classification prediction: Predict the next step as a classification problem (% interval ranges).

In the section Usage you will find how to use the model. Dataset and normalization techniques are covered in the section Dataset. In the section Model you will find the details of the model architectures for each approach. Training details are covered in the section Training (e.g. learning rate functions, performance metrics, etc). The section Results will cover some results obtained with each model and end up with overall conclusions. In the section API you will find the API documentation to better understand how code is structured.

Render of a transformer model as a hologram, projecting from a digital device, with a faint BTC logo in the holographic projection, without any text (DALLE·3)

Contents

• Home
• Usage
  • Installation
  • Running the code
    • Dataset
    • Training
    • Configuration
    • Inference
• Dataset
  • Cryptocurrency data
  • Normalization techniques
  • Real Dataset class
  • Real Dataset configuration
  • Synthetic Dataset class
  • Synthetic Dataset configuration
• Model
  • Backbone
  • Mean prediction
    • Loss Function
  • Distribution prediction
    • Loss Function
  • Classification prediction
    • Loss Function
  • Model Configuration
• Training
  • Configuration
  • Metrics
  • Hyperparameter search
• Results
  • Hardware
  • Mean Prediction
  • Distribution Prediction
  • Classification Prediction
  • Conclusions
• API
  • Modules Overview
  • Run
    • AgentBase
    • FlaxAgent
    • FlaxTrainer
    • TrainerBase
  • Utils
    • Dataset
    • EarlyStopper
    • SyntheticDataset
    • compute_metrics()
    • get_best_model()
    • get_logger()
    • jax_collate_fn()
    • plot_metrics()
    • plot_predictions()
    • plot_tensorboard_experiment()
  • Config
    • DatasetConfig
    • ExperimentConfig
    • ModelConfig
    • SyntheticDatasetConfig
  • Models
    • AddPositionalEncoding
    • Encoder
    • EncoderBlock
    • FeatureExtractor
    • FeedForwardBlock
    • PredictionHead
    • ResidualBlock
    • Time2Vec
    • Transformer
    • TransformerConfig

---

Citation

If you find jaxer useful, please consider citing:

@misc{2024rsanchezmo,
  title     = {Jaxer: Jax and Flax Time Series Prediction Transformer},
  author    = {Rodrigo Sánchez Molina},
  year      = {2024},
  howpublished = {https://github.com/rsanchezmo/jaxer}
}