def tpu_subgraph():
      return tf.tpu.rewrite(model_fn, args)

  return tf.raw_ops.TPUPartitionedCall(
      args=tpu_subgraph.captured_inputs,
      Tout=[o.type for o in tpu_subgraph.definition.signature.output_arg],
      f=tpu_subgraph,
      device_ordinal=[0])

static_shape = tf.placeholder(tf.int32, shape=([3]), name='static_size')

w = tf.Variable(tf.constant([[1.0], [2.0], [3.0]]), name='w')