support multivaraite

Author: Linh-nk

orion/primitives/timesfm.py

@@ -54,27 +54,30 @@ def predict(self, X):
         Args:
             X (ndarray):
                 input timeseries.
-            target (int): 
+            target (int):
                 index of target column in multivariate case. Default to 0.
         Return:
             ndarray:
                 forecasted timeseries.
         """
         frequency_input = [self.freq for _ in range(len(X))]
-        d = X.shape[-1] #number of variables
-        if d == 1: #univariate
+        d = X.shape[-1]
+
+
+        #univariate
+        if d == 1:
             y_hat, _ = self.model.forecast(X[:, :, 0], freq=frequency_input)
             return y_hat[:, 0]
-            
-        else: #multivariate
-            #Extend the x_reg to future values
+
+
+        #multivariate
+        else:
             X_reg = X[:, :, 1:d] 
             m, n, k = X_reg.shape
             X_reg_new = np.zeros((m, n+1, k))
             X_reg_new[:, :-1, :] = X_reg
             X_reg_new[:-1, -1, :] = X_reg[1:, 0, :]
 
-
             x_reg = {str(i): X_reg_new[:, :, i] for i in range(k)}
             y_hat, _ = self.model.forecast_with_covariates(
                 inputs=X[:, :, 0],