bidirectional wrapper for rnns args layer keras layers rnn instance such as keras layers lstm or keras layers gru it could also be a keras layers layer instance that meets the following criteria 1 be a sequenceprocessing layer accepts 3d inputs 2 have a gobackwards returnsequences and returnstate attribute with the same semantics as for the rnn class 3 have an inputspec attribute 4 implement serialization via getconfig and fromconfig note that the recommended way to create new rnn layers is to write a custom rnn cell and use it with keras layers rnn instead of subclassing keras layers layer directly when returnsequences is true the output of the masked timestep will be zero regardless of the layer s original zerooutputformask value mergemode mode by which outputs of the forward and backward rnns will be combined one of sum mul concat ave none if none the outputs will not be combined they will be returned as a list defaults to concat backwardlayer optional keras layers rnn or keras layers layer instance to be used to handle backwards input processing if backwardlayer is not provided the layer instance passed as the layer argument will be used to generate the backward layer automatically note that the provided backwardlayer layer should have properties matching those of the layer argument in particular it should have the same values for stateful returnstates returnsequences etc in addition backwardlayer and layer should have different gobackwards argument values a valueerror will be raised if these requirements are not met call arguments the call arguments for this layer are the same as those of the wrapped rnn layer beware that when passing the initialstate argument during the call of this layer the first half in the list of elements in the initialstate list will be passed to the forward rnn call and the last half in the list of elements will be passed to the backward rnn call note instantiating a bidirectional layer from an existing rnn layer instance will not reuse the weights state of the rnn layer instance the bidirectional layer will have freshly initialized weights examples python model sequential inputshape5 10 bidirectionallstm10 returnsequencestrue bidirectionallstm10 dense5 activationsoftmax model compileloss categoricalcrossentropy optimizer rmsprop with custom backward layer forwardlayer lstm10 returnsequencestrue backwardlayer lstm10 activation relu returnsequencestrue gobackwardstrue model sequential inputshape5 10 bidirectionalforwardlayer backwardlayerbackwardlayer dense5 activationsoftmax model compileloss categoricalcrossentropy optimizer rmsprop recreate the forward layer from the original layer config so that it will not carry over any state from the layer force the zerooutputformask to be true if returning sequences ensure the forward and backward layers have valid common property if self forwardlayer gobackwards self backwardlayer gobackwards raise valueerror forward layer and backward layer should have different gobackwards value received forwardlayer gobackwards fself forwardlayer gobackwards backwardlayer gobackwards fself backwardlayer gobackwards commonattributes stateful returnsequences returnstate for a in commonattributes forwardvalue getattrself forwardlayer a backwardvalue getattrself backwardlayer a if forwardvalue backwardvalue raise valueerror forward layer and backward layer are expected to have f the same value for attribute a got f forwardvalue for forward layer and f backwardvalue for backward layer def computeoutputshapeself sequencesshape initialstateshapenone outputshape self forwardlayer computeoutputshapesequencesshape if self returnstate outputshape stateshape outputshape0 outputshape1 if self mergemode concat outputshape listoutputshape outputshape1 2 outputshape tupleoutputshape elif self mergemode is none outputshape outputshape copy copyoutputshape if self returnstate if self mergemode is none return outputshape stateshape copy copystateshape return outputshape stateshape copy copystateshape return outputshape def call self sequences initialstatenone masknone trainingnone kwargs if self forwardlayer callhastrainingarg kwargstraining training if self forwardlayer callhasmaskarg kwargsmask mask if initialstate is not none initialstates are not keras tensors eg eager tensor from np array they are only passed in from kwarg initialstate and should be passed to forwardbackward layer via kwarg initialstate as well forwardinputs backwardinputs sequences sequences half leninitialstate 2 forwardstate initialstate half backwardstate initialstatehalf else forwardinputs backwardinputs sequences sequences forwardstate backwardstate none none y self forwardlayer forwardinputs initialstateforwardstate kwargs yrev self backwardlayer backwardinputs initialstatebackwardstate kwargs if self returnstate states tupley1 yrev1 y y0 yrev yrev0 y ops casty self computedtype yrev ops castyrev self computedtype if self returnsequences yrev ops flipyrev axis1 if self mergemode concat output ops concatenatey yrev axis1 elif self mergemode sum output y yrev elif self mergemode ave output y yrev 2 elif self mergemode mul output y yrev elif self mergemode is none output y yrev else raise valueerror unrecognized value for mergemode freceived self mergemode expected one of concat sum ave mul if self returnstate if self mergemode is none return output states return output states return output def resetstatesself compatibility alias self resetstate def resetstateself if not self stateful raise attributeerrorlayer must be stateful self forwardlayer resetstate self backwardlayer resetstate property def statesself if self forwardlayer states and self backwardlayer states return tupleself forwardlayer states self backwardlayer states return none def buildself sequencesshape initialstateshapenone self forwardlayer buildsequencesshape self backwardlayer buildsequencesshape self built true def computemaskself mask if isinstancemask list mask mask0 if self returnsequences if not self mergemode outputmask mask mask else outputmask mask else outputmask none none if not self mergemode else none if self returnstate and self states is not none statemask none for in self states if isinstanceoutputmask list return outputmask statemask 2 return outputmask statemask 2 return outputmask def getconfigself config mergemode self mergemode configlayer serializationlib serializekerasobject self forwardlayer configbackwardlayer serializationlib serializekerasobject self backwardlayer baseconfig super getconfig return baseconfig config classmethod def fromconfigcls config customobjectsnone instead of updating the input create a copy and use that config copy deepcopyconfig configlayer serializationlib deserializekerasobject configlayer customobjectscustomobjects handle optional backward layer instantiation backwardlayerconfig config popbackwardlayer none if backwardlayerconfig is not none backwardlayer serializationlib deserializekerasobject backwardlayerconfig customobjectscustomobjects configbackwardlayer backwardlayer instantiate the wrapper adjust it and return it layer clsconfig return layer bidirectional wrapper for rnns args layer keras layers rnn instance such as keras layers lstm or keras layers gru it could also be a keras layers layer instance that meets the following criteria 1 be a sequence processing layer accepts 3d inputs 2 have a go_backwards return_sequences and return_state attribute with the same semantics as for the rnn class 3 have an input_spec attribute 4 implement serialization via get_config and from_config note that the recommended way to create new rnn layers is to write a custom rnn cell and use it with keras layers rnn instead of subclassing keras layers layer directly when return_sequences is true the output of the masked timestep will be zero regardless of the layer s original zero_output_for_mask value merge_mode mode by which outputs of the forward and backward rnns will be combined one of sum mul concat ave none if none the outputs will not be combined they will be returned as a list defaults to concat backward_layer optional keras layers rnn or keras layers layer instance to be used to handle backwards input processing if backward_layer is not provided the layer instance passed as the layer argument will be used to generate the backward layer automatically note that the provided backward_layer layer should have properties matching those of the layer argument in particular it should have the same values for stateful return_states return_sequences etc in addition backward_layer and layer should have different go_backwards argument values a valueerror will be raised if these requirements are not met call arguments the call arguments for this layer are the same as those of the wrapped rnn layer beware that when passing the initial_state argument during the call of this layer the first half in the list of elements in the initial_state list will be passed to the forward rnn call and the last half in the list of elements will be passed to the backward rnn call note instantiating a bidirectional layer from an existing rnn layer instance will not reuse the weights state of the rnn layer instance the bidirectional layer will have freshly initialized weights examples python model sequential input shape 5 10 bidirectional lstm 10 return_sequences true bidirectional lstm 10 dense 5 activation softmax model compile loss categorical_crossentropy optimizer rmsprop with custom backward layer forward_layer lstm 10 return_sequences true backward_layer lstm 10 activation relu return_sequences true go_backwards true model sequential input shape 5 10 bidirectional forward_layer backward_layer backward_layer dense 5 activation softmax model compile loss categorical_crossentropy optimizer rmsprop recreate the forward layer from the original layer config so that it will not carry over any state from the layer force the zero_output_for_mask to be true if returning sequences ensure the forward and backward layers have valid common property initial_states are not keras tensors eg eager tensor from np array they are only passed in from kwarg initial_state and should be passed to forward backward layer via kwarg initial_state as well compatibility alias instead of updating the input create a copy and use that handle optional backward layer instantiation instantiate the wrapper adjust it and return it
import copy

from keras import ops
from keras import utils
from keras.api_export import keras_export
from keras.layers.core.wrapper import Wrapper
from keras.layers.layer import Layer
from keras.saving import serialization_lib


@keras_export("keras.layers.Bidirectional")
class Bidirectional(Wrapper):
    

    def __init__(
        self,
        layer,
        merge_mode="concat",
        weights=None,
        backward_layer=None,
        **kwargs,
    ):
        if not isinstance(layer, Layer):
            raise ValueError(
                "Please initialize `Bidirectional` layer with a "
                f"`keras.layers.Layer` instance. Received: {layer}"
            )
        if backward_layer is not None and not isinstance(backward_layer, Layer):
            raise ValueError(
                "`backward_layer` need to be a `keras.layers.Layer` "
                f"instance. Received: {backward_layer}"
            )
        if merge_mode not in ["sum", "mul", "ave", "concat", None]:
            raise ValueError(
                f"Invalid merge mode. Received: {merge_mode}. "
                "Merge mode should be one of "
                '{"sum", "mul", "ave", "concat", None}'
            )
        super().__init__(layer, **kwargs)

        
        
        config = serialization_lib.serialize_keras_object(layer)
        config["config"]["name"] = "forward_" + utils.removeprefix(
            layer.name, "forward_"
        )
        self.forward_layer = serialization_lib.deserialize_keras_object(config)

        if backward_layer is None:
            config = serialization_lib.serialize_keras_object(layer)
            config["config"]["go_backwards"] = True
            config["config"]["name"] = "backward_" + utils.removeprefix(
                layer.name, "backward_"
            )
            self.backward_layer = serialization_lib.deserialize_keras_object(
                config
            )
        else:
            self.backward_layer = backward_layer
        self._verify_layer_config()

        def force_zero_output_for_mask(layer):
            
            if getattr(layer, "zero_output_for_mask", None) is not None:
                layer.zero_output_for_mask = layer.return_sequences

        force_zero_output_for_mask(self.forward_layer)
        force_zero_output_for_mask(self.backward_layer)

        self.merge_mode = merge_mode
        if weights:
            nw = len(weights)
            self.forward_layer.initial_weights = weights[: nw // 2]
            self.backward_layer.initial_weights = weights[nw // 2 :]
        self.stateful = layer.stateful
        self.return_sequences = layer.return_sequences
        self.return_state = layer.return_state
        self.supports_masking = True
        self.input_spec = layer.input_spec

    def _verify_layer_config(self):
        
        if self.forward_layer.go_backwards == self.backward_layer.go_backwards:
            raise ValueError(
                "Forward layer and backward layer should have different "
                "`go_backwards` value. Received: "
                "forward_layer.go_backwards "
                f"{self.forward_layer.go_backwards}, "
                "backward_layer.go_backwards="
                f"{self.backward_layer.go_backwards}"
            )

        common_attributes = ("stateful", "return_sequences", "return_state")
        for a in common_attributes:
            forward_value = getattr(self.forward_layer, a)
            backward_value = getattr(self.backward_layer, a)
            if forward_value != backward_value:
                raise ValueError(
                    "Forward layer and backward layer are expected to have "
                    f'the same value for attribute "{a}", got '
                    f'"{forward_value}" for forward layer and '
                    f'"{backward_value}" for backward layer'
                )

    def compute_output_shape(self, sequences_shape, initial_state_shape=None):
        output_shape = self.forward_layer.compute_output_shape(sequences_shape)

        if self.return_state:
            output_shape, state_shape = output_shape[0], output_shape[1:]

        if self.merge_mode == "concat":
            output_shape = list(output_shape)
            output_shape[-1] *= 2
            output_shape = tuple(output_shape)
        elif self.merge_mode is None:
            output_shape = [output_shape, copy.copy(output_shape)]

        if self.return_state:
            if self.merge_mode is None:
                return output_shape + state_shape + copy.copy(state_shape)
            return [output_shape] + state_shape + copy.copy(state_shape)
        return output_shape

    def call(
        self,
        sequences,
        initial_state=None,
        mask=None,
        training=None,
    ):
        kwargs = {}
        if self.forward_layer._call_has_training_arg:
            kwargs["training"] = training
        if self.forward_layer._call_has_mask_arg:
            kwargs["mask"] = mask

        if initial_state is not None:
            
            
            
            
            forward_inputs, backward_inputs = sequences, sequences
            half = len(initial_state) // 2
            forward_state = initial_state[:half]
            backward_state = initial_state[half:]
        else:
            forward_inputs, backward_inputs = sequences, sequences
            forward_state, backward_state = None, None

        y = self.forward_layer(
            forward_inputs, initial_state=forward_state, **kwargs
        )
        y_rev = self.backward_layer(
            backward_inputs, initial_state=backward_state, **kwargs
        )

        if self.return_state:
            states = tuple(y[1:] + y_rev[1:])
            y = y[0]
            y_rev = y_rev[0]

        y = ops.cast(y, self.compute_dtype)
        y_rev = ops.cast(y_rev, self.compute_dtype)

        if self.return_sequences:
            y_rev = ops.flip(y_rev, axis=1)
        if self.merge_mode == "concat":
            output = ops.concatenate([y, y_rev], axis=-1)
        elif self.merge_mode == "sum":
            output = y + y_rev
        elif self.merge_mode == "ave":
            output = (y + y_rev) / 2
        elif self.merge_mode == "mul":
            output = y * y_rev
        elif self.merge_mode is None:
            output = (y, y_rev)
        else:
            raise ValueError(
                "Unrecognized value for `merge_mode`. "
                f"Received: {self.merge_mode}"
                'Expected one of {"concat", "sum", "ave", "mul"}.'
            )
        if self.return_state:
            if self.merge_mode is None:
                return output + states
            return (output,) + states
        return output

    def reset_states(self):
        
        self.reset_state()

    def reset_state(self):
        if not self.stateful:
            raise AttributeError("Layer must be stateful.")
        self.forward_layer.reset_state()
        self.backward_layer.reset_state()

    @property
    def states(self):
        if self.forward_layer.states and self.backward_layer.states:
            return tuple(self.forward_layer.states + self.backward_layer.states)
        return None

    def build(self, sequences_shape, initial_state_shape=None):
        self.forward_layer.build(sequences_shape)
        self.backward_layer.build(sequences_shape)
        self.built = True

    def compute_mask(self, _, mask):
        if isinstance(mask, list):
            mask = mask[0]
        if self.return_sequences:
            if not self.merge_mode:
                output_mask = (mask, mask)
            else:
                output_mask = mask
        else:
            output_mask = (None, None) if not self.merge_mode else None

        if self.return_state and self.states is not None:
            state_mask = (None for _ in self.states)
            if isinstance(output_mask, list):
                return output_mask + state_mask * 2
            return (output_mask,) + state_mask * 2
        return output_mask

    def get_config(self):
        config = {"merge_mode": self.merge_mode}
        config["layer"] = serialization_lib.serialize_keras_object(
            self.forward_layer
        )
        config["backward_layer"] = serialization_lib.serialize_keras_object(
            self.backward_layer
        )
        base_config = super().get_config()
        return {**base_config, **config}

    @classmethod
    def from_config(cls, config, custom_objects=None):
        
        config = copy.deepcopy(config)

        config["layer"] = serialization_lib.deserialize_keras_object(
            config["layer"], custom_objects=custom_objects
        )
        
        backward_layer_config = config.pop("backward_layer", None)
        if backward_layer_config is not None:
            backward_layer = serialization_lib.deserialize_keras_object(
                backward_layer_config, custom_objects=custom_objects
            )
            config["backward_layer"] = backward_layer
        
        layer = cls(**config)
        return layer