[Core|Common] TypeError: Got unsupported ScalarType BFloat16

[carusyte]

What happened + What you expected to happen

Nowadays BFloat16 support are getting prevalent for ML libs and hardware. Much appreciated if bfloat16 support can be applied in Neuralforecast as well.

...
831   File "/home/kemove/git/neuralforecast/neuralforecast/core.py", line 1345, in predict_insample
832     model_fcsts = model.predict(trimmed_dataset, step_size=step_size)
833                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
834   File "/home/kemove/git/neuralforecast/neuralforecast/common/_base_multivariate.py", line 607, in predict
835     fcsts = torch.vstack(fcsts).numpy()
836             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
837 TypeError: Got unsupported ScalarType BFloat16
...

Versions / Dependencies

Neuralforecast: 2.0.0

Reproduction script

N/A

Issue Severity

None

[elephaint]

Hey, do you have a piece of code so that I reproduce the error?

[carusyte]

Sure. Generally the error is thrown when I ran this "minimal" reproducer on my workstation with CPU / GPU:

import pandas as pd
import numpy as np
import logging
import torch
from neuralforecast.core import NeuralForecast
from neuralforecast.models import TSMixerx
 
# Prep dummy data
start_date = "2000-01-01"
end_date = "2025-02-01"
date_range = pd.date_range(start=start_date, end=end_date, freq="B")
np.random.seed(0)
df = pd.DataFrame(
    {
        "unique_id": "dummy",
        "ds": date_range,
        "y": np.random.randn(len(date_range)),
        "val1": np.random.randn(len(date_range)),
        "val2": np.random.rand(len(date_range)),
    }
)
logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
torch.set_float32_matmul_precision("medium")
horizon = 10
input_size = 20
models = [
    TSMixerx(
        h=horizon,
        input_size=input_size,
        n_series=1,
        revin=False,
        max_steps=100,
        random_seed=0,
        hist_exog_list=["val1", "val2"],
        precision="bf16-mixed",
    ),
]
nf = NeuralForecast(
    models=models,
    freq="B",
    local_scaler_type="robust",
)
nf.fit(df=df, val_size=300)
Y_hat_insample = nf.predict_insample(step_size=horizon)

However in main branch it's still being blocked by another unresolved issue #1056. I applied a temp fix to side-step that (mentioned here #1056 (comment) ), or you might reuse this branch

Stacktrace:

Traceback (most recent call last):
  File "/home/kemove/Downloads/test.py", line 50, in <module>
    Y_hat_insample = nf.predict_insample(step_size=horizon)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kemove/git/neuralforecast/neuralforecast/core.py", line 1345, in predict_insample
    model_fcsts = model.predict(trimmed_dataset, step_size=step_size)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kemove/git/neuralforecast/neuralforecast/common/_base_multivariate.py", line 607, in predict
    fcsts = torch.vstack(fcsts).numpy()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: Got unsupported ScalarType BFloat16

[elephaint]

Thanks, I can reproduce. If I look at the underlying PyTorch issue I think it makes most sense to just go for the fix mentioned there rather than introducing another package with dependencies for this edge case.

Also the proposed fix converts to uint first, which fails on signed floats.

I've opened a new PR that addresses the issue generically.