(self, xs)
return _make_singleitem_or_original(stack)
Executes this layer as part of a forward pass through the model.

    Args:
      xs: Tensors of as required by the branches of this conditional.

    Returns:
      Tensors resulting from running the chosen branch.
Executes this layer as part of a forward pass through the model.
Executes this layer as part of a forward pass through the model .
def forward(self, xs):
    """Executes this layer as part of a forward pass through the model.

    Args:
      xs: Tensors of as required by the branches of this conditional.

    Returns:
      Tensors resulting from running the chosen branch.
    """
    # TODO(jaszczur): modify; it's a copy from SkippingSerial
    self._validate_forward_inputs(xs)
    layers_state = self.state
    # Get 3 rngs, one for each layer.
    rngs = _split_rngs(self.rng, 3)

    # Prepare the stack and do some safety checks as in the parent class.
    stack = _make_tuple(xs)
    weights = self.weights
    if len(weights) != 3:
      raise ValueError('number of weights ({}) not equal to 3'
                       .format(len(weights)))
    if len(layers_state) != 3:
      raise ValueError('length of state ({}) not equal to 3'
                       .format(len(layers_state)))

    def true_func(t):
      outputs, new_true_state = self._true.pure_fn(
          t[0][0], t[1][0], t[2][0], t[3][0])
      # t[2][1] is old_false_state which is not changing if true is executed.
      return outputs, (new_true_state, t[2][1])

    def false_func(t):
      if self._identity_false_fun:
        # Memory optimization: we don't need pure_fn call.
        return t[0][1], t[2]
      outputs, new_false_state = self._false.pure_fn(
          t[0][1], t[1][1], t[2][1], t[3][1])
      # t[2][1] is old_true_state, which is not changing if false is executed.
      return outputs, (t[2][0], new_false_state)

    cond_inputs = inputs_from_stack(xs, self._cond.n_in)
    cond_output, s = self._cond.pure_fn(cond_inputs, self.weights[0],
                                        self.state[0], rngs[0], use_cache=True)
    stack = outputs_onto_stack([], stack, self._cond.n_in)
    self._cond.state = s

    outputs, both_states = fastmath.cond(
        cond_output,
        true_func,
        false_func,
        [(stack, stack),
         (self.weights[1], self.weights[2]),
         (self.state[1], self.state[2]),
         (rngs[1], rngs[2])]
    )
    stack = outputs_onto_stack([], stack, self._cond.n_in)

    # We don't know which (`true` or `false`) branch was run, but both of them
    # are adding (n_out) and removing (n_in) the same number of elements of the
    # stack (this was checked in __init__). outputs_onto_stack just uses the
    # layer's n_in, so we can pass either `true` or `false` to it.
    # Note that `outputs` is the actual output of `true` or `false` branch,
    # whichever was run, and we add it to the stack in any case.
    stack = outputs_onto_stack(outputs, stack, self._true.n_in)
    self._true.state = both_states[0]
    self._false.state = both_states[1]
    return _make_singleitem_or_original(stack)
def forward ( self , xs ) : # TODO(jaszczur): modify; it's a copy from SkippingSerial self . _validate_forward_inputs ( xs ) layers_state = self . state # Get 3 rngs, one for each layer. rngs = _split_rngs ( self . rng , 3 ) # Prepare the stack and do some safety checks as in the parent class. stack = _make_tuple ( xs ) weights = self . weights if len ( weights ) != 3 : raise ValueError ( 'number of weights ({}) not equal to 3' . format ( len ( weights ) ) ) if len ( layers_state ) != 3 : raise ValueError ( 'length of state ({}) not equal to 3' . format ( len ( layers_state ) ) ) def true_func ( t ) : outputs , new_true_state = self . _true . pure_fn ( t [ 0 ] [ 0 ] , t [ 1 ] [ 0 ] , t [ 2 ] [ 0 ] , t [ 3 ] [ 0 ] ) # t[2][1] is old_false_state which is not changing if true is executed. return outputs , ( new_true_state , t [ 2 ] [ 1 ] ) def false_func ( t ) : if self . _identity_false_fun : # Memory optimization: we don't need pure_fn call. return t [ 0 ] [ 1 ] , t [ 2 ] outputs , new_false_state = self . _false . pure_fn ( t [ 0 ] [ 1 ] , t [ 1 ] [ 1 ] , t [ 2 ] [ 1 ] , t [ 3 ] [ 1 ] ) # t[2][1] is old_true_state, which is not changing if false is executed. return outputs , ( t [ 2 ] [ 0 ] , new_false_state ) cond_inputs = inputs_from_stack ( xs , self . _cond . n_in ) cond_output , s = self . _cond . pure_fn ( cond_inputs , self . weights [ 0 ] , self . state [ 0 ] , rngs [ 0 ] , use_cache = True ) stack = outputs_onto_stack ( [ ] , stack , self . _cond . n_in ) self . _cond . state = s outputs , both_states = fastmath . cond ( cond_output , true_func , false_func , [ ( stack , stack ) , ( self . weights [ 1 ] , self . weights [ 2 ] ) , ( self . state [ 1 ] , self . state [ 2 ] ) , ( rngs [ 1 ] , rngs [ 2 ] ) ] ) stack = outputs_onto_stack ( [ ] , stack , self . _cond . n_in ) # We don't know which (`true` or `false`) branch was run, but both of them # are adding (n_out) and removing (n_in) the same number of elements of the # stack (this was checked in __init__). outputs_onto_stack just uses the # layer's n_in, so we can pass either `true` or `false` to it. # Note that `outputs` is the actual output of `true` or `false` branch, # whichever was run, and we add it to the stack in any case. stack = outputs_onto_stack ( outputs , stack , self . _true . n_in ) self . _true . state = both_states [ 0 ] self . _false . state = both_states [ 1 ] return _make_singleitem_or_original ( stack )