rasa-pro 3.14.0a20__py3-none-any.whl → 3.14.0a23__py3-none-any.whl

rasa/utils/tensorflow/callback.py

@@ -2,111 +2,146 @@ import logging
 from pathlib import Path
 from typing import Any, Dict, Optional, Text
 
-import tensorflow as tf
-from tqdm import tqdm
+from rasa.utils.tensorflow import TENSORFLOW_AVAILABLE
+
+if TENSORFLOW_AVAILABLE:
+    import tensorflow as tf
+    from tqdm import tqdm
+else:
+    # Placeholder values when TensorFlow is not available
+    tf = None
+    tqdm = None
 
 import rasa.shared.utils.io
 
 logger = logging.getLogger(__name__)
 
 
-class RasaTrainingLogger(tf.keras.callbacks.Callback):
-    """Callback for logging the status of training."""
-
-    def __init__(self, epochs: int, silent: bool) -> None:
-        """Initializes the callback.
-
-        Args:
-            epochs: Total number of epochs.
-            silent: If 'True' the entire progressbar wrapper is disabled.
-        """
-        super().__init__()
-
-        disable = silent or rasa.shared.utils.io.is_logging_disabled()
-        self.progress_bar = tqdm(range(epochs), desc="Epochs", disable=disable)
-
-    def on_epoch_end(self, epoch: int, logs: Optional[Dict[Text, Any]] = None) -> None:
-        """Updates the logging output on every epoch end.
-
-        Args:
-            epoch: The current epoch.
-            logs: The training metrics.
-        """
-        self.progress_bar.update(1)
-        self.progress_bar.set_postfix(logs)
-
-    def on_train_end(self, logs: Optional[Dict[Text, Any]] = None) -> None:
-        """Closes the progress bar after training.
-
-        Args:
-            logs: The training metrics.
-        """
-        self.progress_bar.close()
-
-
-class RasaModelCheckpoint(tf.keras.callbacks.Callback):
-    """Callback for saving intermediate model checkpoints."""
-
-    def __init__(self, checkpoint_dir: Path) -> None:
-        """Initializes the callback.
-
-        Args:
-            checkpoint_dir: Directory to store checkpoints to.
-        """
-        super().__init__()
-
-        self.checkpoint_file = checkpoint_dir / "checkpoint.tf_model"
-        self.best_metrics_so_far: Dict[Text, Any] = {}
-
-    def on_epoch_end(self, epoch: int, logs: Optional[Dict[Text, Any]] = None) -> None:
-        """Save the model on epoch end if the model has improved.
-
-        Args:
-            epoch: The current epoch.
-            logs: The training metrics.
-        """
-        if self._does_model_improve(logs):
-            logger.debug(f"Creating model checkpoint at epoch={epoch + 1} ...")
-            self.model.save_weights(
-                self.checkpoint_file, overwrite=True, save_format="tf"
-            )
-
-    def _does_model_improve(self, curr_results: Dict[Text, Any]) -> bool:
-        """Checks whether the current results are better than the best so far.
-
-        Results are considered better if each metric is equal or better than the best so
-        far, and at least one is better.
-
-        Args:
-            curr_results: The training metrics for this epoch.
-        """
-        curr_metric_names = [
-            k
-            for k in curr_results.keys()
-            if k.startswith("val") and (k.endswith("_acc") or k.endswith("_f1"))
-        ]
-        # the "val" prefix is prepended to metrics in fit if _should_eval returns true
-        # for this particular epoch
-        if len(curr_metric_names) == 0:
-            # the metrics are not validation metrics
-            return False
-        # initialize best_metrics_so_far with the first results
-        if not self.best_metrics_so_far:
-            for metric_name in curr_metric_names:
-                self.best_metrics_so_far[metric_name] = float(curr_results[metric_name])
-            return True
-
-        at_least_one_improved = False
-        improved_metrics = {}
-        for metric_name in self.best_metrics_so_far.keys():
-            if float(curr_results[metric_name]) < self.best_metrics_so_far[metric_name]:
-                # at least one of the values is worse
+if TENSORFLOW_AVAILABLE:
+
+    class RasaTrainingLogger(tf.keras.callbacks.Callback):
+        """Callback for logging the status of training."""
+
+        def __init__(self, epochs: int, silent: bool) -> None:
+            """Initializes the callback.
+
+            Args:
+                epochs: Total number of epochs.
+                silent: If 'True' the entire progressbar wrapper is disabled.
+            """
+            super().__init__()
+
+            disable = silent or rasa.shared.utils.io.is_logging_disabled()
+            self.progress_bar = tqdm(range(epochs), desc="Epochs", disable=disable)
+
+        def on_epoch_end(
+            self, epoch: int, logs: Optional[Dict[Text, Any]] = None
+        ) -> None:
+            """Updates the logging output on every epoch end.
+
+            Args:
+                epoch: The current epoch.
+                logs: The training metrics.
+            """
+            self.progress_bar.update(1)
+            self.progress_bar.set_postfix(logs)
+
+        def on_train_end(self, logs: Optional[Dict[Text, Any]] = None) -> None:
+            """Closes the progress bar after training.
+
+            Args:
+                logs: The training metrics.
+            """
+            self.progress_bar.close()
+
+    class RasaModelCheckpoint(tf.keras.callbacks.Callback):
+        """Callback for saving intermediate model checkpoints."""
+
+        def __init__(self, checkpoint_dir: Path) -> None:
+            """Initializes the callback.
+
+            Args:
+                checkpoint_dir: Directory to store checkpoints to.
+            """
+            super().__init__()
+
+            self.checkpoint_file = checkpoint_dir / "checkpoint.weights.h5"
+            self.best_metrics_so_far: Dict[Text, Any] = {}
+
+        def on_epoch_end(
+            self, epoch: int, logs: Optional[Dict[Text, Any]] = None
+        ) -> None:
+            """Save the model on epoch end if the model has improved.
+
+            Args:
+                epoch: The current epoch.
+                logs: The training metrics.
+            """
+            if self._does_model_improve(logs):
+                logger.debug(f"Creating model checkpoint at epoch={epoch + 1} ...")
+                # Ensure model is built before saving weights
+                if not self.model.built:
+                    # Build the model with dummy data to ensure it's built
+                    import tensorflow as tf
+
+                    dummy_input = tf.zeros((1, 1))
+                    _ = self.model(dummy_input)
+
+                # Ensure the directory exists before saving
+                import os
+
+                os.makedirs(os.path.dirname(self.checkpoint_file), exist_ok=True)
+                self.model.save_weights(self.checkpoint_file, overwrite=True)
+
+        def _does_model_improve(self, curr_results: Dict[Text, Any]) -> bool:
+            """Checks whether the current results are better than the best so far.
+
+            Results are considered better if each metric is
+            equal or better than the best so far, and at least one is better.
+
+            Args:
+                curr_results: The training metrics for this epoch.
+            """
+            curr_metric_names = [
+                k
+                for k in curr_results.keys()
+                if k.startswith("val") and (k.endswith("_acc") or k.endswith("_f1"))
+            ]
+            # the "val" prefix is prepended to metrics in fit
+            # if _should_eval returns true
+            # for this particular epoch
+            if len(curr_metric_names) == 0:
+                # the metrics are not validation metrics
                 return False
-            if float(curr_results[metric_name]) > self.best_metrics_so_far[metric_name]:
-                at_least_one_improved = True
-                improved_metrics[metric_name] = float(curr_results[metric_name])
-
-        # all current values >= previous best and at least one is better
-        if at_least_one_improved:
-            self.best_metrics_so_far.update(improved_metrics)
-        return at_least_one_improved
+            # initialize best_metrics_so_far with the first results
+            if not self.best_metrics_so_far:
+                for metric_name in curr_metric_names:
+                    self.best_metrics_so_far[metric_name] = float(
+                        curr_results[metric_name]
+                    )
+                    return True
+
+            at_least_one_improved = False
+            improved_metrics = {}
+            for metric_name in self.best_metrics_so_far.keys():
+                if (
+                    float(curr_results[metric_name])
+                    < self.best_metrics_so_far[metric_name]
+                ):
+                    # at least one of the values is worse
+                    return False
+                if (
+                    float(curr_results[metric_name])
+                    > self.best_metrics_so_far[metric_name]
+                ):
+                    at_least_one_improved = True
+                    improved_metrics[metric_name] = float(curr_results[metric_name])
+
+            # all current values >= previous best and at least one is better
+            if at_least_one_improved:
+                self.best_metrics_so_far.update(improved_metrics)
+            return at_least_one_improved
+else:
+    # Placeholder classes when TensorFlow is not available
+    RasaTrainingLogger = None  # type: ignore
+    RasaModelCheckpoint = None  # type: ignore

rasa/utils/tensorflow/crf.py

@@ -9,7 +9,7 @@ from tensorflow.types.experimental import TensorLike
 # (modified to our neeeds)
 
 
-class CrfDecodeForwardRnnCell(tf.keras.layers.AbstractRNNCell):
+class CrfDecodeForwardRnnCell(tf.keras.layers.Layer):
     """Computes the forward decoding in a linear-chain CRF."""
 
     def __init__(self, transition_params: TensorLike, **kwargs: Any) -> None:

rasa/utils/tensorflow/data_generator.py

@@ -71,13 +71,22 @@ class RasaDataGenerator(Sequence):
         # balancing on the next epoch
         return data
 
+    @staticmethod
+    def _create_default_array() -> np.ndarray:
+        """Create a default empty array for missing features.
+
+        Returns:
+            A default empty array with shape (0, 1) and dtype float32.
+        """
+        return np.zeros((0, 1), dtype=np.float32)
+
     @staticmethod
     def prepare_batch(
         data: Data,
         start: Optional[int] = None,
         end: Optional[int] = None,
         tuple_sizes: Optional[Dict[Text, int]] = None,
-    ) -> Tuple[Optional[np.ndarray], ...]:
+    ) -> Tuple[np.ndarray, ...]:
         """Slices model data into batch using given start and end value.
 
         Args:
@@ -85,8 +94,8 @@ class RasaDataGenerator(Sequence):
             start: The start index of the batch
             end: The end index of the batch
             tuple_sizes: In case the feature is not present we propagate the batch with
-              None. Tuple sizes contains the number of how many None values to add for
-              what kind of feature.
+              default arrays. Tuple sizes contains the number of how many default values
+              to add for what kind of feature.
 
         Returns:
             The features of the batch.
@@ -95,12 +104,14 @@ class RasaDataGenerator(Sequence):
 
         for key, attribute_data in data.items():
             for sub_key, f_data in attribute_data.items():
-                # add None for not present values during processing
+                # add default arrays for not present values during processing
                 if not f_data:
                     if tuple_sizes:
-                        batch_data += [None] * tuple_sizes[key]
+                        batch_data += [
+                            RasaDataGenerator._create_default_array()
+                        ] * tuple_sizes[key]
                     else:
-                        batch_data.append(None)
+                        batch_data.append(RasaDataGenerator._create_default_array())
                     continue
 
                 for v in f_data:
@@ -409,8 +420,10 @@ class RasaBatchDataGenerator(RasaDataGenerator):
         end = start + self._current_batch_size
 
         # return input and target data, as our target data is inside the input
-        # data return None for the target data
-        return self.prepare_batch(self._data, start, end), None
+        # data return default array for the target data
+        return self.prepare_batch(
+            self._data, start, end
+        ), RasaDataGenerator._create_default_array()
 
     def on_epoch_end(self) -> None:
         """Update the data after every epoch."""

rasa/utils/tensorflow/layers.py

@@ -3,9 +3,7 @@ from typing import Any, Callable, List, Optional, Text, Tuple, Union
 
 import tensorflow as tf
 import tensorflow.keras.backend as K
-
-# TODO: The following is not (yet) available via tf.keras
-from keras.src.utils.control_flow_util import smart_cond
+from tensorflow.python.keras.utils.control_flow_util import smart_cond
 
 import rasa.utils.tensorflow.crf
 import rasa.utils.tensorflow.layers_utils as layers_utils
@@ -278,6 +276,7 @@ class RandomlyConnectedDense(tf.keras.layers.Dense):
             kernel_constraint: Constraint function applied to
                 the `kernel` weights matrix.
             bias_constraint: Constraint function applied to the bias vector.
+            **kwargs: Additional keyword arguments passed to the parent class.
         """
         super().__init__(**kwargs)
 
@@ -298,16 +297,19 @@ class RandomlyConnectedDense(tf.keras.layers.Dense):
             self.kernel_mask = None
             return
 
-        # Construct mask with given density and guarantee that every output is
-        # connected to at least one input
-        kernel_mask = self._minimal_mask() + self._random_mask()
+        # Use callable initializer for TensorFlow 2.19.1 compatibility
+        def kernel_mask_initializer() -> tf.Tensor:
+            # Construct mask with given density and guarantee that every output is
+            # connected to at least one input
+            kernel_mask = self._minimal_mask() + self._random_mask()
 
-        # We might accidently have added a random connection on top of
-        # a fixed connection
-        kernel_mask = tf.clip_by_value(kernel_mask, 0, 1)
+            # We might accidently have added a random connection on top of
+            # a fixed connection
+            kernel_mask = tf.clip_by_value(kernel_mask, 0, 1)
+            return kernel_mask
 
         self.kernel_mask = tf.Variable(
-            initial_value=kernel_mask, trainable=False, name="kernel_mask"
+            initial_value=kernel_mask_initializer, trainable=False, name="kernel_mask"
         )
 
     def _random_mask(self) -> tf.Tensor:
@@ -367,7 +369,12 @@ class RandomlyConnectedDense(tf.keras.layers.Dense):
         Returns:
             The processed inputs.
         """
-        if self.density < 1.0:
+        # Apply kernel masking if needed (Keras 3.x compatibility check)
+        if (
+            self.density < 1.0
+            and hasattr(self, "kernel_mask")
+            and self.kernel_mask is not None
+        ):
             # Set fraction of the `kernel` weights to zero according to precomputed mask
             self.kernel.assign(self.kernel * self.kernel_mask)
         return super().call(inputs)
@@ -724,6 +731,7 @@ class DotProductLoss(tf.keras.layers.Layer):
                 Currently, the only possible value is `SOFTMAX`.
             similarity_type: Similarity measure to use, either `cosine` or `inner`.
             name: Optional name of the layer.
+            **kwargs: Additional keyword arguments passed to the parent class.
 
         Raises:
             TFLayerConfigException: When `similarity_type` is not one of `COSINE` or
@@ -883,6 +891,7 @@ class SingleLabelDotProductLoss(DotProductLoss):
                 values are approximately bounded.
             model_confidence: Normalization of confidence values during inference.
                 Currently, the only possible value is `SOFTMAX`.
+            **kwargs: Additional keyword arguments passed to the parent class.
         """
         super().__init__(
             num_candidates,
@@ -1244,6 +1253,7 @@ class MultiLabelDotProductLoss(DotProductLoss):
                 Used inside _loss_cross_entropy() only.
             model_confidence: Normalization of confidence values during inference.
                 Currently, the only possible value is `SOFTMAX`.
+            **kwargs: Additional keyword arguments passed to the parent class.
         """
         super().__init__(
             num_candidates,

rasa/utils/tensorflow/metrics.py

@@ -1,7 +1,6 @@
 from typing import Any, Dict, Optional
 
 import tensorflow as tf
-from tensorflow.keras import backend as K
 from tensorflow.types.experimental import TensorLike
 
 # original code taken from
@@ -118,7 +117,7 @@ class FBetaScore(tf.keras.metrics.Metric):
 
         def _zero_wt_init(name: Any) -> Any:
             return self.add_weight(
-                name, shape=self.init_shape, initializer="zeros", dtype=self.dtype
+                name=name, shape=self.init_shape, initializer="zeros", dtype=self.dtype
             )
 
         self.true_positives = _zero_wt_init("true_positives")
@@ -197,7 +196,12 @@ class FBetaScore(tf.keras.metrics.Metric):
 
     def reset_state(self) -> None:
         reset_value = tf.zeros(self.init_shape, dtype=self.dtype)
-        K.batch_set_value([(v, reset_value) for v in self.variables])
+        # In Keras 3.x, self.variables contains string names, not variable objects,
+        # so each metric variable is reset using assign() instead of K.batch_set_value()
+        self.true_positives.assign(reset_value)
+        self.false_positives.assign(reset_value)
+        self.false_negatives.assign(reset_value)
+        self.weights_intermediate.assign(reset_value)
 
     def reset_states(self) -> None:
         # Backwards compatibility alias of `reset_state`. New classes should

rasa/utils/tensorflow/models.py

@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text, Tuple, Union
 import numpy as np
 import tensorflow as tf
 from keras import Model
-from keras.src.utils import tf_utils
 
 import rasa.utils.train_utils
 from rasa.shared.constants import DIAGNOSTIC_DATA
@@ -78,6 +77,7 @@ class RasaModel(Model):
 
         Args:
             random_seed: set the random seed to get reproducible results
+            **kwargs: Additional keyword arguments passed to the parent class
         """
         # make sure that keras releases resources from previously trained model
         tf.keras.backend.clear_session()
@@ -273,7 +273,8 @@ class RasaModel(Model):
         if self._run_eagerly:
             # Once we take advantage of TF's distributed training, this is where
             # scheduled functions will be forced to execute and return actual values.
-            outputs = tf_utils.sync_to_numpy_or_python_type(self.predict_step(batch_in))
+            step_output = self.predict_step(batch_in)
+            outputs = self._convert_tensors_to_numpy(step_output)
             if DIAGNOSTIC_DATA in outputs:
                 outputs[DIAGNOSTIC_DATA] = self._empty_lists_to_none_in_dict(
                     outputs[DIAGNOSTIC_DATA]
@@ -287,9 +288,8 @@ class RasaModel(Model):
 
         # Once we take advantage of TF's distributed training, this is where
         # scheduled functions will be forced to execute and return actual values.
-        outputs = tf_utils.sync_to_numpy_or_python_type(
-            self._tf_predict_step(list(batch_in))
-        )
+        step_output = self._tf_predict_step(list(batch_in))
+        outputs = self._convert_tensors_to_numpy(step_output)
         if DIAGNOSTIC_DATA in outputs:
             outputs[DIAGNOSTIC_DATA] = self._empty_lists_to_none_in_dict(
                 outputs[DIAGNOSTIC_DATA]
@@ -388,6 +388,35 @@ class RasaModel(Model):
 
         return {k: _recurse(v) for k, v in input_dict.items()}
 
+    def _convert_tensors_to_numpy(
+        self, step_output: Dict[Text, Any]
+    ) -> Dict[Text, Any]:
+        """ "Recursively convert TensorFlow tensors to numpy arrays for Keras 3.x
+        compatibility.
+        Replaces the deprecated tf_utils.sync_to_numpy_or_python_type() function.
+        Converts tensors (objects with 'numpy' method) to numpy arrays,
+        leaves others unchanged.
+        """
+
+        def to_numpy(obj: Any) -> Any:
+            if hasattr(obj, "numpy"):
+                try:
+                    return obj.numpy()
+                except Exception:
+                    # Fallback: return as-is if not convertible
+                    return obj
+            # Dict: recurse into values
+            if isinstance(obj, dict):
+                return {k: to_numpy(v) for k, v in obj.items()}
+            # List/Tuple: recurse preserving type
+            if isinstance(obj, list):
+                return [to_numpy(v) for v in obj]
+            if isinstance(obj, tuple):
+                return tuple(to_numpy(v) for v in obj)
+            return obj
+
+        return {key: to_numpy(value) for key, value in step_output.items()}
+
     def _get_metric_results(self, prefix: Optional[Text] = "") -> Dict[Text, float]:
         return {
             f"{prefix}{metric.name}": metric.result()
@@ -403,7 +432,21 @@ class RasaModel(Model):
             overwrite: If 'True' an already existing model with the same file name will
                        be overwritten.
         """
-        self.save_weights(model_file_name, overwrite=overwrite, save_format="tf")
+        # Ensure filename ends with .weights.h5 and model is built for Keras 3.x
+        # compatibility
+        model_file_name = str(model_file_name)
+        if not model_file_name.endswith(".weights.h5"):
+            model_file_name += ".weights.h5"
+
+        if not self.built:
+            import tensorflow as tf
+
+            _ = self(tf.zeros((1, 1)))
+
+        # TensorFlow 2.19: save weights with different file extension
+        if not model_file_name.endswith(".weights.h5"):
+            model_file_name += ".weights.h5"
+        self.save_weights(model_file_name, overwrite=overwrite)
 
     @classmethod
     def load(
@@ -438,18 +481,23 @@ class RasaModel(Model):
         learning_rate = kwargs.get("config", {}).get(LEARNING_RATE, 0.001)
         run_eagerly = kwargs.get("config", {}).get(RUN_EAGERLY)
 
-        # need to train on 1 example to build weights of the correct size
         model.compile(
             optimizer=tf.keras.optimizers.Adam(learning_rate), run_eagerly=run_eagerly
         )
         data_generator = RasaBatchDataGenerator(model_data_example, batch_size=1)
         model.fit(data_generator, verbose=False)
+
+        # Ensure model is built before loading weights
+        if not model.built:
+            sample_batch = next(iter(data_generator))
+            _ = model(sample_batch)
+
         # load trained weights
         model.load_weights(model_file_name)
 
         # predict on one data example to speed up prediction during inference
         # the first prediction always takes a bit longer to trace tf function
-        if not finetune_mode and predict_data_example:
+        if predict_data_example:
             model.run_inference(predict_data_example)
 
         logger.debug("Finished loading the model.")

rasa/utils/tensorflow/rasa_layers.py

@@ -291,7 +291,7 @@ class ConcatenateSparseDenseFeatures(RasaCustomLayer):
             [
                 config[DENSE_DIMENSION][attribute]
                 if signature.is_sparse
-                else signature.units
+                else (signature.units or 0)
                 for signature in feature_type_signature
             ]
         )
@@ -301,12 +301,12 @@ class ConcatenateSparseDenseFeatures(RasaCustomLayer):
     ) -> tf.Tensor:
         """Turns sparse tensor into dense, possibly adds dropout before and/or after."""
         if self.SPARSE_DROPOUT in self._tf_layers:
-            feature = self._tf_layers[self.SPARSE_DROPOUT](feature, training)
+            feature = self._tf_layers[self.SPARSE_DROPOUT](feature, training=training)
 
         feature = self._tf_layers[self.SPARSE_TO_DENSE](feature)
 
         if self.DENSE_DROPOUT in self._tf_layers:
-            feature = self._tf_layers[self.DENSE_DROPOUT](feature, training)
+            feature = self._tf_layers[self.DENSE_DROPOUT](feature, training=training)
 
         return feature
 
@@ -944,7 +944,7 @@ class RasaSequenceLayer(RasaCustomLayer):
         # Note that only sequence-level features are masked, nothing happens to the
         # sentence-level features in the combined features tensor.
         seq_sent_features, mlm_boolean_mask = self._tf_layers[self.MLM_INPUT_MASK](
-            seq_sent_features, mask_sequence, training
+            seq_sent_features, mask_sequence, training=training
         )
 
         return seq_sent_features, token_ids, mlm_boolean_mask
@@ -1002,7 +1002,9 @@ class RasaSequenceLayer(RasaCustomLayer):
         ]((sequence_features, sentence_features, sequence_feature_lengths))
 
         # Apply one or more dense layers.
-        seq_sent_features = self._tf_layers[self.FFNN](seq_sent_features, training)
+        seq_sent_features = self._tf_layers[self.FFNN](
+            seq_sent_features, training=training
+        )
 
         # If using masked language modeling, mask the transformer inputs and get labels
         # for the masked tokens and a boolean mask. Note that TED does not use MLM loss,
@@ -1031,7 +1033,7 @@ class RasaSequenceLayer(RasaCustomLayer):
         if self._has_transformer:
             mask_padding = 1 - mask_combined_sequence_sentence
             outputs, attention_weights = self._tf_layers[self.TRANSFORMER](
-                seq_sent_features_masked, mask_padding, training
+                seq_sent_features_masked, mask_padding, training=training
             )
             outputs = tf.nn.gelu(outputs)
         else:

rasa/utils/tensorflow/transformer.py

@@ -2,10 +2,8 @@ from typing import Optional, Text, Tuple, Union
 
 import numpy as np
 import tensorflow as tf
-
-# TODO: The following is not (yet) available via tf.keras
-from keras.src.utils.control_flow_util import smart_cond
 from tensorflow.keras import backend as K
+from tensorflow.python.keras.utils.control_flow_util import smart_cond
 
 from rasa.utils.tensorflow.exceptions import TFLayerConfigException
 from rasa.utils.tensorflow.layers import RandomlyConnectedDense
@@ -280,6 +278,7 @@ class MultiHeadAttention(tf.keras.layers.Layer):
             value: A tensor with shape (..., length, depth).
             pad_mask: Float tensor with shape broadcastable
                 to (..., length, length). Defaults to None.
+            training: A tensor
 
         Returns:
             output: A tensor with shape (..., length, depth).