tf-keras-nightly 2.17.0.dev2024050509__py3-none-any.whl → 2.19.0.dev2024101709__py3-none-any.whl

tf_keras/__init__.py

@@ -27,4 +27,4 @@ from tf_keras.src.engine.sequential import Sequential
 from tf_keras.src.engine.training import Model
 
 
-__version__ = "2.17.0.dev2024050509"
+__version__ = "2.19.0.dev2024101709"

tf_keras/src/__init__.py

@@ -35,7 +35,7 @@ from tf_keras.src.testing_infra import test_utils
 from tensorflow.python import tf2
 from tensorflow.python.util.tf_export import keras_export
 
-__version__ = "2.17.0"
+__version__ = "2.19.0"
 
 keras_export("keras.__version__").export_constant(__name__, "__version__")
 

tf_keras/src/callbacks.py

@@ -1423,20 +1423,20 @@ class ModelCheckpoint(Callback):
         if mode == "min":
             self.monitor_op = np.less
             if self.best is None:
-                self.best = np.Inf
+                self.best = np.inf
         elif mode == "max":
             self.monitor_op = np.greater
             if self.best is None:
-                self.best = -np.Inf
+                self.best = -np.inf
         else:
             if "acc" in self.monitor or self.monitor.startswith("fmeasure"):
                 self.monitor_op = np.greater
                 if self.best is None:
-                    self.best = -np.Inf
+                    self.best = -np.inf
             else:
                 self.monitor_op = np.less
                 if self.best is None:
-                    self.best = np.Inf
+                    self.best = np.inf
 
         if self.save_freq != "epoch" and not isinstance(self.save_freq, int):
             raise ValueError(
@@ -1903,6 +1903,23 @@ class BackupAndRestore(Callback):
                 "only supports empty strategy, "
                 "MirroredStrategy, MultiWorkerMirroredStrategy and TPUStrategy."
             )
+
+        # Re-initialize the optimizer.
+        if self.model.built:
+            if (
+                self.model.optimizer is not None
+                and callable(getattr(self.model.optimizer, "build", None))
+                and not getattr(self.model.optimizer, "_built", False)
+            ):
+                self.model.optimizer.build(self.model.trainable_variables)
+        else:
+            logging.warning(
+                "To use the BackupAndRestore callback, "
+                "you model must be built before you call `fit()`. "
+                f"Model {self.model} is unbuilt. You can build it "
+                "beforehand by calling it on a batch of data."
+            )
+
         self.model._training_state = worker_training_state.WorkerTrainingState(
             self.model,
             self.backup_dir,
@@ -2095,7 +2112,7 @@ class EarlyStopping(Callback):
         # Allow instances to be re-used
         self.wait = 0
         self.stopped_epoch = 0
-        self.best = np.Inf if self.monitor_op == np.less else -np.Inf
+        self.best = np.inf if self.monitor_op == np.less else -np.inf
         self.best_weights = None
         self.best_epoch = 0
 
@@ -3098,10 +3115,10 @@ class ReduceLROnPlateau(Callback):
             self.mode == "auto" and "acc" not in self.monitor
         ):
             self.monitor_op = lambda a, b: np.less(a, b - self.min_delta)
-            self.best = np.Inf
+            self.best = np.inf
         else:
             self.monitor_op = lambda a, b: np.greater(a, b + self.min_delta)
-            self.best = -np.Inf
+            self.best = -np.inf
         self.cooldown_counter = 0
         self.wait = 0
 

tf_keras/src/engine/base_layer.py

@@ -578,7 +578,8 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
             Accepted values are constants defined in the class
             `tf.VariableAggregation`.
           **kwargs: Additional keyword arguments. Accepted values are `getter`,
-            `collections`, `experimental_autocast` and `caching_device`.
+            `collections`, `autocast`, `experimental_autocast` and
+            `caching_device`.
 
         Returns:
           The variable created.
@@ -594,6 +595,7 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
         # Validate optional keyword arguments.
         for kwarg in kwargs:
             if kwarg not in [
+                "autocast",
                 "collections",
                 "experimental_autocast",
                 "caching_device",
@@ -603,9 +605,13 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
             ]:
                 raise TypeError("Unknown keyword argument:", kwarg)
         collections_arg = kwargs.pop("collections", None)
-        # 'experimental_autocast' can be set to False by the caller to indicate
-        # an AutoCastVariable should never be created.
-        autocast = kwargs.pop("experimental_autocast", True)
+        # 'autocast' or 'experimental_autocast' can be set to False by the
+        # caller to indicate an AutoCastVariable should never be created.
+        autocast = kwargs.pop("autocast", None)
+        if autocast is None:
+            autocast = kwargs.pop("experimental_autocast", None)
+            if autocast is None:
+                autocast = True
         # See the docstring for tf.Variable about the details for
         # caching_device.
         caching_device = kwargs.pop("caching_device", None)

tf_keras/src/engine/base_layer_v1.py

@@ -352,7 +352,8 @@ class Layer(base_layer.Layer):
             Accepted values are constants defined in the class
             `tf.VariableAggregation`.
           **kwargs: Additional keyword arguments. Accepted values are `getter`,
-            `collections`, `experimental_autocast` and `caching_device`.
+            `collections`, `autocast`, `experimental_autocast` and
+            `caching_device`.
 
         Returns:
           The created variable. Usually either a `Variable` or
@@ -371,6 +372,7 @@ class Layer(base_layer.Layer):
         # Validate optional keyword arguments.
         for kwarg in kwargs:
             if kwarg not in [
+                "autocast",
                 "getter",
                 "collections",
                 "experimental_autocast",
@@ -380,9 +382,13 @@ class Layer(base_layer.Layer):
         has_custom_getter = "getter" in kwargs
         getter = kwargs.pop("getter", base_layer_utils.make_variable)
         collections_arg = kwargs.pop("collections", None)
-        # 'experimental_autocast' can be set to False by the caller to indicate
-        # an AutoCastVariable should never be created.
-        autocast = kwargs.pop("experimental_autocast", True)
+        # 'autocast' or 'experimental_autocast' can be set to False by the
+        # caller to indicate an AutoCastVariable should never be created.
+        autocast = kwargs.pop("autocast", None)
+        if autocast is None:
+            autocast = kwargs.pop("experimental_autocast", None)
+            if autocast is None:
+                autocast = True
         # See the docstring for tf.Variable about the details for
         # caching_device.
         caching_device = kwargs.pop("caching_device", None)

tf_keras/src/engine/node.py

@@ -84,9 +84,10 @@ class Node:
         self.call_args = call_args
         self.call_kwargs = call_kwargs
 
-        # Cached for performance.
+        # Cached for performance. Put kwargs in order of the call method instead
+        # of using the sorted key order from `tf.nest.flatten`.
         self._flat_arguments = tf.nest.flatten(
-            (self.call_args, self.call_kwargs)
+            (self.call_args, self.call_kwargs.values())
         )
         # Used to avoid expensive `nest` operations in the most common case.
         self._single_positional_tensor_passed = (
@@ -176,9 +177,13 @@ class Node:
             for kt_id, kt_index in self._keras_inputs_ids_and_indices:
                 flat_arguments[kt_index] = tensor_dict[kt_id].pop()
 
+            # Pack the same way as `self._flat_arguments`, i.e. `kwargs` as a
+            # list in the original order.
             args, kwargs = tf.nest.pack_sequence_as(
-                (self.call_args, self.call_kwargs), flat_arguments
+                (self.call_args, self.call_kwargs.values()), flat_arguments
             )
+            # Add the keys to `kwargs` to go from a list to a dict.
+            kwargs = {k: v for k, v in zip(self.call_kwargs.keys(), kwargs)}
             return args, kwargs
 
     def serialize(self, make_node_key, node_conversion_map):

tf_keras/src/layers/activation/prelu.py

@@ -102,7 +102,7 @@ class PReLU(Layer):
                 if i not in self.shared_axes:
                     axes[i] = input_shape[i]
         self.input_spec = InputSpec(ndim=len(input_shape), axes=axes)
-        self.built = True
+        super().build(input_shape)
 
     def call(self, inputs):
         pos = backend.relu(inputs)

tf_keras/src/layers/attention/base_dense_attention.py

@@ -86,7 +86,8 @@ class BaseDenseAttention(base_layer.BaseRandomLayer):
         # be purely stateless, with no reference to any variable.
         if self.dropout > 0:
             super().build(input_shape)
-        self.built = True
+        else:
+            base_layer.Layer.build(self, input_shape)
 
     def _calculate_scores(self, query, key):
         """Calculates attention scores.

tf_keras/src/layers/convolutional/base_conv.py

@@ -248,7 +248,7 @@ class Conv(Layer):
         self.input_spec = InputSpec(
             min_ndim=self.rank + 2, axes={channel_axis: input_channel}
         )
-        self.built = True
+        super().build(input_shape)
 
     def convolution_op(self, inputs, kernel):
         if self.padding == "causal":

tf_keras/src/layers/convolutional/base_depthwise_conv.py

@@ -20,6 +20,7 @@ import tensorflow.compat.v2 as tf
 from tf_keras.src import constraints
 from tf_keras.src import initializers
 from tf_keras.src import regularizers
+from tf_keras.src.engine.base_layer import Layer
 from tf_keras.src.engine.input_spec import InputSpec
 from tf_keras.src.layers.convolutional.base_conv import Conv
 
@@ -202,7 +203,8 @@ class DepthwiseConv(Conv):
         self.input_spec = InputSpec(
             min_ndim=self.rank + 2, axes={channel_axis: input_dim}
         )
-        self.built = True
+        # Call Layer.build() to skip Conv.build() which we override here.
+        Layer.build(self, input_shape)
 
     def call(self, inputs):
         raise NotImplementedError

tf_keras/src/layers/convolutional/base_separable_conv.py

@@ -21,6 +21,7 @@ from tf_keras.src import activations
 from tf_keras.src import constraints
 from tf_keras.src import initializers
 from tf_keras.src import regularizers
+from tf_keras.src.engine.base_layer import Layer
 from tf_keras.src.engine.input_spec import InputSpec
 from tf_keras.src.layers.convolutional.base_conv import Conv
 
@@ -203,7 +204,8 @@ class SeparableConv(Conv):
             )
         else:
             self.bias = None
-        self.built = True
+        # Call Layer.build() to skip Conv.build() which we override here.
+        Layer.build(self, input_shape)
 
     def call(self, inputs):
         raise NotImplementedError

tf_keras/src/layers/convolutional/conv1d_transpose.py

@@ -22,6 +22,7 @@ from tf_keras.src import constraints
 from tf_keras.src import initializers
 from tf_keras.src import regularizers
 from tf_keras.src.dtensor import utils
+from tf_keras.src.engine.base_layer import Layer
 from tf_keras.src.engine.input_spec import InputSpec
 from tf_keras.src.layers.convolutional.conv1d import Conv1D
 from tf_keras.src.utils import conv_utils
@@ -214,7 +215,8 @@ class Conv1DTranspose(Conv1D):
             )
         else:
             self.bias = None
-        self.built = True
+        # Call Layer.build() to skip Conv.build() which we override here.
+        Layer.build(self, input_shape)
 
     def call(self, inputs):
         inputs_shape = tf.shape(inputs)

tf_keras/src/layers/convolutional/conv2d_transpose.py

@@ -23,6 +23,7 @@ from tf_keras.src import constraints
 from tf_keras.src import initializers
 from tf_keras.src import regularizers
 from tf_keras.src.dtensor import utils
+from tf_keras.src.engine.base_layer import Layer
 from tf_keras.src.engine.input_spec import InputSpec
 from tf_keras.src.layers.convolutional.conv2d import Conv2D
 from tf_keras.src.utils import conv_utils
@@ -240,7 +241,8 @@ class Conv2DTranspose(Conv2D):
             )
         else:
             self.bias = None
-        self.built = True
+        # Call Layer.build() to skip Conv.build() which we override here.
+        Layer.build(self, input_shape)
 
     def call(self, inputs):
         inputs_shape = tf.shape(inputs)

tf_keras/src/layers/convolutional/conv3d_transpose.py

@@ -22,6 +22,7 @@ from tf_keras.src import constraints
 from tf_keras.src import initializers
 from tf_keras.src import regularizers
 from tf_keras.src.dtensor import utils
+from tf_keras.src.engine.base_layer import Layer
 from tf_keras.src.engine.input_spec import InputSpec
 from tf_keras.src.layers.convolutional.conv3d import Conv3D
 from tf_keras.src.utils import conv_utils
@@ -247,7 +248,8 @@ class Conv3DTranspose(Conv3D):
             )
         else:
             self.bias = None
-        self.built = True
+        # Call Layer.build() to skip Conv.build() which we override here.
+        Layer.build(self, input_shape)
 
     def call(self, inputs):
         inputs_shape = tf.shape(inputs)

tf_keras/src/layers/core/dense.py

@@ -174,7 +174,7 @@ class Dense(Layer):
             )
         else:
             self.bias = None
-        self.built = True
+        super().build(input_shape)
 
     def call(self, inputs):
         if inputs.dtype.base_dtype != self._compute_dtype_object.base_dtype:

tf_keras/src/layers/core/embedding.py

@@ -185,7 +185,7 @@ class Embedding(Layer):
             constraint=self.embeddings_constraint,
             experimental_autocast=False,
         )
-        self.built = True
+        super().build(input_shape)
 
     def compute_mask(self, inputs, mask=None):
         if not self.mask_zero:

tf_keras/src/layers/locally_connected/locally_connected1d.py

@@ -284,7 +284,7 @@ class LocallyConnected1D(Layer):
             self.input_spec = InputSpec(ndim=3, axes={1: input_dim})
         else:
             self.input_spec = InputSpec(ndim=3, axes={-1: input_dim})
-        self.built = True
+        super().build(input_shape)
 
     @tf_utils.shape_type_conversion
     def compute_output_shape(self, input_shape):

tf_keras/src/layers/locally_connected/locally_connected2d.py

@@ -308,7 +308,7 @@ class LocallyConnected2D(Layer):
             self.input_spec = InputSpec(ndim=4, axes={1: input_filter})
         else:
             self.input_spec = InputSpec(ndim=4, axes={-1: input_filter})
-        self.built = True
+        super().build(input_shape)
 
     @tf_utils.shape_type_conversion
     def compute_output_shape(self, input_shape):

tf_keras/src/layers/normalization/batch_normalization.py

@@ -542,7 +542,7 @@ class BatchNormalizationBase(Layer):
         finally:
             if partitioner:
                 self._scope.set_partitioner(partitioner)
-        self.built = True
+        super().build(input_shape)
 
     def call(self, inputs, training=None, mask=None):
         inputs = tf.cast(inputs, self.compute_dtype)

tf_keras/src/layers/normalization/layer_normalization.py

@@ -249,7 +249,7 @@ class LayerNormalization(Layer):
             self.beta = None
 
         self._fused = self._fused_can_be_used(rank)
-        self.built = True
+        super().build(input_shape)
 
     def call(self, inputs):
         # TODO(b/229545225): Remove the RaggedTensor check.

tf_keras/src/layers/rnn/abstract_rnn_cell.py

@@ -56,7 +56,7 @@ class AbstractRNNCell(base_layer.Layer):
               shape=(self.units, self.units),
               initializer='uniform',
               name='recurrent_kernel')
-          self.built = True
+          super().build(input_shape)
 
         def call(self, inputs, states):
           prev_output = states[0]

tf_keras/src/layers/rnn/base_conv_lstm.py

@@ -218,7 +218,6 @@ class ConvLSTMCell(DropoutRNNCellMixin, base_layer.BaseRandomLayer):
             )
         else:
             self.bias = None
-        self.built = True
 
     def call(self, inputs, states, training=None):
         h_tm1 = states[0]  # previous memory state

tf_keras/src/layers/rnn/base_conv_rnn.py

@@ -20,6 +20,7 @@ import tensorflow.compat.v2 as tf
 
 from tf_keras.src import backend
 from tf_keras.src.engine import base_layer
+from tf_keras.src.engine.base_layer import Layer
 from tf_keras.src.engine.input_spec import InputSpec
 from tf_keras.src.layers.rnn.base_rnn import RNN
 from tf_keras.src.utils import conv_utils
@@ -207,6 +208,8 @@ class ConvRNN(RNN):
 
     @tf_utils.shape_type_conversion
     def build(self, input_shape):
+        # Call Layer.build() to skip RNN.build() which we override here.
+        Layer.build(self, input_shape)
         # Note input_shape will be list of shapes of initial states and
         # constants if these are passed in __call__.
         if self._num_constants is not None:
@@ -263,7 +266,6 @@ class ConvRNN(RNN):
                 ]
         if self.stateful:
             self.reset_states()
-        self.built = True
 
     def get_initial_state(self, inputs):
         # (samples, timesteps, img_dims..., filters)

tf_keras/src/layers/rnn/base_rnn.py

@@ -207,7 +207,7 @@ class RNN(base_layer.Layer):
                 shape=(self.units, self.units),
                 initializer='uniform',
                 name='recurrent_kernel')
-            self.built = True
+            super().build(input_shape)
 
         def call(self, inputs, states):
             prev_output = states[0]

tf_keras/src/layers/rnn/base_wrapper.py

@@ -56,7 +56,7 @@ class Wrapper(Layer):
         if not self.layer.built:
             self.layer.build(input_shape)
             self.layer.built = True
-        self.built = True
+        super().build(input_shape)
 
     @property
     def activity_regularizer(self):

tf_keras/src/layers/rnn/bidirectional.py

@@ -470,7 +470,8 @@ class Bidirectional(Wrapper):
             self.forward_layer.build(input_shape)
         with backend.name_scope(self.backward_layer.name):
             self.backward_layer.build(input_shape)
-        self.built = True
+        # Call Layer.build() to skip Wrapper.build() which we override here.
+        Layer.build(self, input_shape)
 
     def compute_mask(self, inputs, mask):
         if isinstance(mask, list):

tf_keras/src/layers/rnn/cell_wrappers.py

@@ -102,10 +102,10 @@ class _RNNCellWrapper(AbstractRNNCell):
             inputs, state, cell_call_fn=self.cell.call, **kwargs
         )
 
-    def build(self, inputs_shape):
+    def build(self, input_shape):
         """Builds the wrapped cell."""
-        self.cell.build(inputs_shape)
-        self.built = True
+        self.cell.build(input_shape)
+        super().build(input_shape)
 
     @property
     def wrapped_cell(self):

tf_keras/src/layers/rnn/cudnn_gru.py

@@ -144,8 +144,6 @@ class CuDNNGRU(_CuDNNRNN):
             constraint=self.bias_constraint,
         )
 
-        self.built = True
-
     def _process_batch(self, inputs, initial_state):
         if not self.time_major:
             inputs = tf.transpose(inputs, perm=(1, 0, 2))
@@ -172,6 +170,10 @@ class CuDNNGRU(_CuDNNRNN):
             shape=self._vector_shape,
         )
 
+        batch_dim = tf.shape(inputs)[1]
+        max_sequence_length = tf.shape(inputs)[0]
+        sequence_lengths = tf.fill([batch_dim], max_sequence_length)
+
         args = {
             "input": inputs,
             "input_h": input_h,
@@ -179,9 +181,10 @@ class CuDNNGRU(_CuDNNRNN):
             "params": params,
             "is_training": True,
             "rnn_mode": "gru",
+            "sequence_lengths": sequence_lengths,
         }
 
-        outputs, h, _, _, _ = tf.raw_ops.CudnnRNNV2(**args)
+        outputs, h, _, _, _ = tf.raw_ops.CudnnRNNV3(**args)
 
         if self.stateful or self.return_state:
             h = h[0]

tf_keras/src/layers/rnn/cudnn_lstm.py

@@ -170,8 +170,6 @@ class CuDNNLSTM(_CuDNNRNN):
             constraint=self.bias_constraint,
         )
 
-        self.built = True
-
     def _process_batch(self, inputs, initial_state):
         if not self.time_major:
             inputs = tf.transpose(inputs, perm=(1, 0, 2))
@@ -204,15 +202,20 @@ class CuDNNLSTM(_CuDNNRNN):
             shape=self._vector_shape,
         )
 
+        batch_dim = tf.shape(inputs)[1]
+        max_sequence_length = tf.shape(inputs)[0]
+        sequence_lengths = tf.fill([batch_dim], max_sequence_length)
+
         args = {
             "input": inputs,
             "input_h": input_h,
             "input_c": input_c,
             "params": params,
             "is_training": True,
+            "sequence_lengths": sequence_lengths,
         }
 
-        outputs, h, c, _, _ = tf.raw_ops.CudnnRNNV2(**args)
+        outputs, h, c, _, _ = tf.raw_ops.CudnnRNNV3(**args)
 
         if self.stateful or self.return_state:
             h = h[0]

tf_keras/src/layers/rnn/gru.py

@@ -222,7 +222,6 @@ class GRUCell(DropoutRNNCellMixin, base_layer.BaseRandomLayer):
             )
         else:
             self.bias = None
-        self.built = True
 
     def call(self, inputs, states, training=None):
         h_tm1 = (
@@ -1034,11 +1033,13 @@ def gpu_gru(
             mask, time_major
         )
 
-    if not time_major and sequence_lengths is None:
-        inputs = tf.transpose(inputs, perm=(1, 0, 2))
-        seq_axis, batch_axis = (0, 1)
-    else:
-        seq_axis, batch_axis = (0, 1) if time_major else (1, 0)
+    seq_axis, batch_axis = (0, 1) if time_major else (1, 0)
+
+    if sequence_lengths is None:
+        max_sequence_length = tf.shape(inputs)[seq_axis]
+        batch_size = tf.shape(inputs)[batch_axis]
+        sequence_lengths = tf.fill([batch_size], max_sequence_length)
+
     # For init_h, cuDNN expects one more dim of num_layers before or after batch
     # dim for time major or batch major inputs respectively
     init_h = tf.expand_dims(init_h, axis=seq_axis)
@@ -1069,49 +1070,36 @@ def gpu_gru(
         transpose_weights=True,
     )
 
-    if sequence_lengths is not None:
-        if go_backwards:
-            # Three reversals are required. E.g.,
-            # normal input = [1, 2, 3, 0, 0]  # where 0 need to be masked
-            # reversed_input_to_cudnn = [3, 2, 1, 0, 0]
-            # output_from_cudnn = [6, 5, 4, 0, 0]
-            # expected_output = [0, 0, 6, 5 ,4]
-            inputs = tf.reverse_sequence(
-                inputs,
-                sequence_lengths,
-                seq_axis=seq_axis,
-                batch_axis=batch_axis,
-            )
-        outputs, h, _, _, _ = tf.raw_ops.CudnnRNNV3(
-            input=inputs,
-            input_h=init_h,
-            input_c=0,
-            params=params,
-            is_training=True,
-            rnn_mode="gru",
-            sequence_lengths=sequence_lengths,
-            time_major=time_major,
+    if go_backwards:
+        # Three reversals are required. E.g.,
+        # normal input = [1, 2, 3, 0, 0]  # where 0 need to be masked
+        # reversed_input_to_cudnn = [3, 2, 1, 0, 0]
+        # output_from_cudnn = [6, 5, 4, 0, 0]
+        # expected_output = [0, 0, 6, 5 ,4]
+        inputs = tf.reverse_sequence(
+            inputs,
+            sequence_lengths,
+            seq_axis=seq_axis,
+            batch_axis=batch_axis,
         )
-        if go_backwards:
-            outputs = tf.reverse_sequence(
-                outputs,
-                sequence_lengths,
-                seq_axis=seq_axis,
-                batch_axis=batch_axis,
-            )
-            outputs = tf.reverse(outputs, axis=[seq_axis])
-    else:
-        if go_backwards:
-            # Reverse axis 0 since the input is already convert to time major.
-            inputs = tf.reverse(inputs, axis=[0])
-        outputs, h, _, _ = tf.raw_ops.CudnnRNN(
-            input=inputs,
-            input_h=init_h,
-            input_c=0,
-            params=params,
-            is_training=True,
-            rnn_mode="gru",
+    outputs, h, _, _, _ = tf.raw_ops.CudnnRNNV3(
+        input=inputs,
+        input_h=init_h,
+        input_c=0,
+        params=params,
+        is_training=True,
+        rnn_mode="gru",
+        sequence_lengths=sequence_lengths,
+        time_major=time_major,
+    )
+    if go_backwards:
+        outputs = tf.reverse_sequence(
+            outputs,
+            sequence_lengths,
+            seq_axis=seq_axis,
+            batch_axis=batch_axis,
         )
+        outputs = tf.reverse(outputs, axis=[seq_axis])
 
     last_output = outputs[-1]
     if not time_major and sequence_lengths is None and return_sequences:

tf_keras/src/layers/rnn/legacy_cell_wrappers.py

@@ -368,9 +368,9 @@ class DropoutWrapper(_RNNCellWrapperV1):
     def wrapped_cell(self):
         return self.cell
 
-    def build(self, inputs_shape):
-        self.cell.build(inputs_shape)
-        self.built = True
+    def build(self, input_shape):
+        self.cell.build(input_shape)
+        super().build(input_shape)
 
     def _variational_recurrent_dropout_value(
         self, unused_index, value, noise, keep_prob