b2756dda944a6a65e98bedc51cbe42238222ccf0,examples/toy_model_tpu.py,,toy_model,#Any#Any#,97

Before Change

  io_dim = mtf.Dimension("io", FLAGS.io_size)

  x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
  h = mtf.layers.dense(x, hidden_dim, name="layer1", use_bias=False)
  y = mtf.layers.dense(h, io_dim, name="layer2", use_bias=False)

  loss = mtf.reduce_sum(mtf.square(y - x))
  return y, loss

After Change

  x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
  x = mtf.cast(x, activation_dtype)
  h = x
  for lnum in xrange(FLAGS.num_hidden_layers + 1):
    if lnum + 1 == FLAGS.num_hidden_layers + 1:
      dim = io_dim
    elif lnum % 2 == 0:
      dim = mtf.Dimension("hidden_even", FLAGS.hidden_size)
    else:
      dim = mtf.Dimension("hidden_odd", FLAGS.hidden_size)
    h = mtf.layers.dense(
        h, dim,
        use_bias=False,
        master_dtype=master_dtype,
        slice_dtype=slice_dtype,
        name="layer_%d" % lnum)
  y = h

  loss = mtf.reduce_mean(mtf.square(y - x))
  return y, loss

Instances