tensorflow_probability/examples/grammar_vae.py
def call(self, inputs):
    """Runs the model forward to generate a sequence of productions.

    Args:
      inputs: Unused.

    Returns:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.
    """
    del inputs  # unused
    latent_code = ed.MultivariateNormalDiag(loc=tf.zeros(self.latent_size),
                                            sample_shape=1,
                                            name="latent_code")
    state = self.lstm.zero_state(1, dtype=tf.float32)
    t = 0
    productions = []
    stack = [self.grammar.start_symbol]
    while stack:
      symbol = stack.pop()
      net, state = self.lstm(latent_code, state)
      logits = (self.output_layer(net) +
                self.grammar.mask(symbol, on_value=0., off_value=-1e9))
      production = ed.OneHotCategorical(logits=logits,
                                        name="production_" + str(t))
      _, rhs = self.grammar.production_rules[tf.argmax(
          input=production, axis=-1)]
      for symbol in rhs:
        if symbol in self.grammar.nonterminal_symbols:
          stack.append(symbol)
      productions.append(production)
      t += 1
    return tf.stack(productions, axis=1)
def call(self, inputs):
    """Runs the model forward to generate a sequence of productions.

    Args:
      inputs: Unused.

    Returns:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.
    """
    del inputs  # unused
    latent_code = ed.MultivariateNormalDiag(loc=tf.zeros(self.latent_size),
                                            sample_shape=1,
                                            name="latent_code")
    state = self.lstm.zero_state(1, dtype=tf.float32)
    t = 0
    productions = []
    stack = [self.grammar.start_symbol]
    while stack:
      symbol = stack.pop()
      net, state = self.lstm(latent_code, state)
      logits = (self.output_layer(net) +
                self.grammar.mask(symbol, on_value=0., off_value=-1e9))
      production = ed.OneHotCategorical(logits=logits,
                                        name="production_" + str(t))
      _, rhs = self.grammar.production_rules[tf.argmax(
          input=production, axis=-1)]
      for symbol in rhs:
        if symbol in self.grammar.nonterminal_symbols:
          stack.append(symbol)
      productions.append(production)
      t += 1
    return tf.stack(productions, axis=1)
Runs the model forward to generate a sequence of productions.

    Args:
      inputs: Unused.

    Returns:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.