tensorflow_probability/python/bijectors/masked_autoregressive.py
def build(self, input_shape):
    """See tfkl.Layer.build."""
    if self._event_shape is None:
      # `event_shape` wasn't specied at __init__, so infer from `input_shape`.
      self._event_shape = [tf.compat.dimension_value(input_shape[-1])]
      self._event_size = self._event_shape[-1]
      self._event_ndims = len(self._event_shape)
      # Should we throw if input_shape has rank > 2?

    if input_shape[-1] != self._event_shape[-1]:
      raise ValueError("Invalid final dimension of `input_shape`. "
                       "Expected `{!r}`, but got `{!r}`".format(
                           self._event_shape[-1], input_shape[-1]))

    # Construct the masks.
    self._input_order = _create_input_order(
        self._event_size, self._input_order_param)
    self._masks = _create_masks(_create_degrees(
        input_size=self._event_size,
        hidden_units=self._hidden_units,
        input_order=self._input_order,
        hidden_degrees=self._hidden_degrees))

    # In the final layer, we will produce `self._params` outputs for each of the
    # `self._event_size` inputs to `AutoregressiveLayer`.  But `masks[-1]` has
    # shape `[self._hidden_units[-1], self._event_size]`.  Thus, we need to
    # expand the mask to `[hidden_units[-1], event_size * self._params]` such
    # that all units for the same input are masked identically.  In particular,
    # we tile the mask so the j-th element of `tf.unstack(output, axis=-1)` is a
    # tensor of the j-th parameter/unit for each input.
    #
    # NOTE: Other orderings of the output could be faster -- should benchmark.
    self._masks[-1] = np.reshape(
        np.tile(self._masks[-1][..., tf.newaxis], [1, 1, self._params]),
        [self._masks[-1].shape[0], self._event_size * self._params])

    self._network = tf.keras.Sequential([
        # Starting this model with an `InputLayer` ensures that Keras will build
        # and propagate our `dtype` to each layer we add.
        tf.keras.layers.InputLayer((self._event_size,), dtype=self.dtype)
    ])

    # Input-to-hidden, hidden-to-hidden, and hidden-to-output layers:
    #  [..., self._event_size] -> [..., self._hidden_units[0]].
    #  [..., self._hidden_units[k-1]] -> [..., self._hidden_units[k]].
    #  [..., self._hidden_units[-1]] -> [..., event_size * self._params].
    layer_output_sizes = self._hidden_units + [self._event_size * self._params]
    for k in range(len(self._masks)):
      self._network.add(tf.keras.layers.Dense(
          layer_output_sizes[k],
          kernel_initializer=_make_masked_initializer(
              self._masks[k], self._kernel_initializer),
          kernel_constraint=_make_masked_constraint(self._masks[k]),
          activation=self._activation if k + 1 < len(self._masks) else None,
          use_bias=self._use_bias,
          **self._kwargs))

    # Record that the layer has been built.
    super(AutoregressiveLayer, self).build(input_shape)
def build(self, input_shape):
    """See tfkl.Layer.build."""
    if self._event_shape is None:
      # `event_shape` wasn't specied at __init__, so infer from `input_shape`.
      self._event_shape = [tf.compat.dimension_value(input_shape[-1])]
      self._event_size = self._event_shape[-1]
      self._event_ndims = len(self._event_shape)
      # Should we throw if input_shape has rank > 2?

    if input_shape[-1] != self._event_shape[-1]:
      raise ValueError("Invalid final dimension of `input_shape`. "
                       "Expected `{!r}`, but got `{!r}`".format(
                           self._event_shape[-1], input_shape[-1]))

    # Construct the masks.
    self._input_order = _create_input_order(
        self._event_size, self._input_order_param)
    self._masks = _create_masks(_create_degrees(
        input_size=self._event_size,
        hidden_units=self._hidden_units,
        input_order=self._input_order,
        hidden_degrees=self._hidden_degrees))

    # In the final layer, we will produce `self._params` outputs for each of the
    # `self._event_size` inputs to `AutoregressiveLayer`.  But `masks[-1]` has
    # shape `[self._hidden_units[-1], self._event_size]`.  Thus, we need to
    # expand the mask to `[hidden_units[-1], event_size * self._params]` such
    # that all units for the same input are masked identically.  In particular,
    # we tile the mask so the j-th element of `tf.unstack(output, axis=-1)` is a
    # tensor of the j-th parameter/unit for each input.
    #
    # NOTE: Other orderings of the output could be faster -- should benchmark.
    self._masks[-1] = np.reshape(
        np.tile(self._masks[-1][..., tf.newaxis], [1, 1, self._params]),
        [self._masks[-1].shape[0], self._event_size * self._params])

    self._network = tf.keras.Sequential([
        # Starting this model with an `InputLayer` ensures that Keras will build
        # and propagate our `dtype` to each layer we add.
        tf.keras.layers.InputLayer((self._event_size,), dtype=self.dtype)
    ])

    # Input-to-hidden, hidden-to-hidden, and hidden-to-output layers:
    #  [..., self._event_size] -> [..., self._hidden_units[0]].
    #  [..., self._hidden_units[k-1]] -> [..., self._hidden_units[k]].
    #  [..., self._hidden_units[-1]] -> [..., event_size * self._params].
    layer_output_sizes = self._hidden_units + [self._event_size * self._params]
    for k in range(len(self._masks)):
      self._network.add(tf.keras.layers.Dense(
          layer_output_sizes[k],
          kernel_initializer=_make_masked_initializer(
              self._masks[k], self._kernel_initializer),
          kernel_constraint=_make_masked_constraint(self._masks[k]),
          activation=self._activation if k + 1 < len(self._masks) else None,
          use_bias=self._use_bias,
          **self._kwargs))

    # Record that the layer has been built.
    super(AutoregressiveLayer, self).build(input_shape)
See tfkl.Layer.build.