def make_grid_context(n):
  """Create a KG encoding a 2-d grid with four+ relations."""

  # set up the context object
  if FLAGS.variant == 'mix':
    tf.logging.info('using MixtureNQExpr')
    context = nql.NeuralQueryContext()
    context.expression_factory_class = MixtureNQExpr
  elif FLAGS.variant == 'sum':
    tf.logging.info('using MatSumNQExpr')
    context = nql.NeuralQueryContext()
    context.expression_factory_class = MatSumNQExpr
    if FLAGS.minibatch_size != 1:
      raise NotImplementedError('sum not implemented for minibatch_size!=1')
  else:
    context = nql.NeuralQueryContext()
    tf.logging.info('using join/std follow operation')

  # declare the basic directions
  context.declare_relation('n', 'place_t', 'place_t')
  context.declare_relation('s', 'place_t', 'place_t')
  context.declare_relation('e', 'place_t', 'place_t')
  context.declare_relation('w', 'place_t', 'place_t')

  # add in any extra relations
  extra_rels = [('r%d' % i) for i in range(FLAGS.num_extra_rels)]
  tf.logging.info('extra_rels %r' % extra_rels[-10:])
  for r in extra_rels:
    context.declare_relation(r, 'place_t', 'place_t')

  # build a list of facts to put in the KB, one fact per line in format
  #    rel <TAB> subj <TAB> obj
  kg_lines = []
  dij = {'n': (-1, 0), 's': (+1, 0), 'e': (0, +1), 'w': (0, -1)}
  t0 = time.time()
  tf.logging.info('creating kg lines...')
  for i in range(0, n):
    for j in range(0, n):
      for direction, (di, dj) in dij.items():
        if extra_rels:
          direction = extra_rels.pop()
        if (0 <= i + di < n) and (0 <= j + dj < n):
          kg_lines.append(
              '\t'.join([direction, cell(i, j),
                         cell(i + di, j + dj)]) + '\n')

  tf.logging.info('created lines in %f sec' % (time.time() - t0))
  tf.logging.info('%d extra rels remain of %d' %
                  (len(extra_rels), FLAGS.num_extra_rels))
  tf.logging.info('loading %d lines' % len(kg_lines))
  for i in range(min(20, len(kg_lines))):
    tf.logging.info('kg_line %02d: %s' % ((i + 1), kg_lines[i].strip()))

  # load the facts into the KB
  t0 = time.time()
  context.load_kg(lines=kg_lines)
  tf.logging.info('loaded in %f sec' % (time.time() - t0))

  # define the KB type for relations
  context.construct_relation_group('dir_g', 'place_t', 'place_t')

  # return the context holding the KB
  return context
Create a KG encoding a 2-d grid with four+ relations.
language/nql/demos/gridworld_scaling/scaling_eval.py