UnicoLab
diff --git a/‎kdp/custom_layers.py‎
Lines changed: 46 additions & 16 deletions b/‎kdp/custom_layers.py‎
Lines changed: 46 additions & 16 deletions
@@ -104,12 +104,16 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
             inputs (tf.Tensor): Tensor with input data.
 
         Returns:
-            tf.Tensor: processed date tensor with all components [year, month, day_of_month, day_of_week].
+            tf.Tensor: processed date tensor with all components
+            [year, month, day_of_month, day_of_week].
         """
 
         def parse_date(date_str: str) -> tf.Tensor:
             # Handle missing/invalid dates
-            is_valid = tf.strings.regex_full_match(date_str, r"^\d{1,4}[-/]\d{1,2}[-/]\d{1,2}$")
+            is_valid = tf.strings.regex_full_match(
+                date_str,
+                r"^\d{1,4}[-/]\d{1,2}[-/]\d{1,2}$",
+            )
             tf.debugging.assert_equal(
                 is_valid,
                 True,
@@ -126,24 +130,48 @@ def parse_date(date_str: str) -> tf.Tensor:
 
             # Validate date components
             # Validate year is in reasonable range
-            tf.debugging.assert_greater_equal(year, 1000, message="Year must be >= 1000")
-            tf.debugging.assert_less_equal(year, 2200, message="Year must be <= 2200")
+            tf.debugging.assert_greater_equal(
+                year,
+                1000,
+                message="Year must be >= 1000",
+            )
+            tf.debugging.assert_less_equal(
+                year,
+                2200,
+                message="Year must be <= 2200",
+            )
 
             # Validate month is between 1-12
-            tf.debugging.assert_greater_equal(month, 1, message="Month must be >= 1")
-            tf.debugging.assert_less_equal(month, 12, message="Month must be <= 12")
+            tf.debugging.assert_greater_equal(
+                month,
+                1,
+                message="Month must be >= 1",
+            )
+            tf.debugging.assert_less_equal(
+                month,
+                12,
+                message="Month must be <= 12",
+            )
 
             # Validate day is between 1-31
-            tf.debugging.assert_greater_equal(day_of_month, 1, message="Day must be >= 1")
-            tf.debugging.assert_less_equal(day_of_month, 31, message="Day must be <= 31")
+            tf.debugging.assert_greater_equal(
+                day_of_month,
+                1,
+                message="Day must be >= 1",
+            )
+            tf.debugging.assert_less_equal(
+                day_of_month,
+                31,
+                message="Day must be <= 31",
+            )
 
             # Calculate day of week using Zeller's congruence
             y = tf.where(month < 3, year - 1, year)
             m = tf.where(month < 3, month + 12, month)
             k = y % 100
             j = y // 100
             h = (day_of_month + ((13 * (m + 1)) // 5) + k + (k // 4) + (j // 4) - (2 * j)) % 7
-            day_of_week = tf.where(h == 0, 6, h - 1)  # Adjust to 0-6 range where 0 is Monday
+            day_of_week = tf.where(h == 0, 6, h - 1)  # Adjust to 0-6 range where 0 is Sunday
 
             return tf.stack([year, month, day_of_month, day_of_week])
 
@@ -268,10 +296,11 @@ class SeasonLayer(tf.keras.layers.Layer):
     Spring, Summer, and Fall. The one-hot encoding is appended to the input tensor.
 
     Required Input Format:
-        - A tensor of shape [batch_size, 3], where each row contains:
+        - A tensor of shape [batch_size, 4], where each row contains:
             - year (int): Year as a numerical value.
             - month (int): Month as an integer from 1 to 12.
-            - day_of_week (int): Day of the week as an integer from 0 to 6 (where 0=Monday).
+            - day_of_month (int): Day of the month as an integer from 1 to 31.
+            - day_of_week (int): Day of the week as an integer from 0 to 6 (where 0=Sunday).
     """
 
     def __init__(self, **kwargs):
@@ -283,14 +312,15 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
         """Adds seasonal one-hot encoding to the input tensor.
 
         Args:
-            inputs (tf.Tensor): A tensor of shape [batch_size, 3] where each row contains [year, month, day_of_week].
+            inputs (tf.Tensor): A tensor of shape [batch_size, 4] where each row contains
+            [year, month, day_of_month, day_of_week].
 
         Returns:
-            tf.Tensor: A tensor of shape [batch_size, 7] with the original features
+            tf.Tensor: A tensor of shape [batch_size, 8] with the original features
             plus the one-hot encoded season information.
 
         Raises:
-            ValueError: If the input tensor does not have shape [batch_size, 3] or contains invalid month values.
+            ValueError: If the input tensor does not have shape [batch_size, 4] or contains invalid month values.
         """
         # Ensure inputs is 2D
         if len(tf.shape(inputs)) == 1:
@@ -312,8 +342,8 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
             + tf.cast(is_fall, tf.int32) * 3
         )
 
-        # Convert season to one-hot encoding
-        season_one_hot = tf.one_hot(season, depth=4)
+        # Convert season to one-hot encoding and cast to int32 to match input type
+        season_one_hot = tf.cast(tf.one_hot(season, depth=4), tf.int32)
 
         return tf.concat([inputs, season_one_hot], axis=-1)