def FProp(self, theta, inputs):
"""Apply projection to inputs.
Args:
theta: A NestedMap object containing weights' values of this layer and its
children layers.
inputs: The inputs tensor. Shaped [..., input_dims].
Returns:
Projected inputs.
"""
p = self.params
with tf.name_scope(p.name):
computation_cost.Add(
self, 'flops',
tf.reduce_prod(tf.to_int64(tf.shape(inputs)[:-1])) * tf.to_int64(
symbolic.EvalExpr(symbolic.TENSOR_VALUES,
p.input_dims * p.output_dims)) * 2)
use_tpu = py_utils.use_tpu()
if use_tpu and inputs.shape is not None and inputs.shape.rank < 26:
# Avoids reshape if feasible and uses Einsum.
if inputs.shape.rank == 2:
return tf.matmul(inputs, theta.w)
else:
s = ''.join([chr(x) for x in range(97, 123)]) # abc...xyz
r = inputs.shape.rank
return tf.einsum('{0}y,yz->{0}z'.format(s[:r - 1]), inputs, theta.w)
input_dim = py_utils.GetShape(inputs)[-1]
act = tf.matmul(tf.reshape(inputs, [-1, input_dim]), theta.w)
output_dim = tf.shape(theta.w)[-1]
act = tf.reshape(act,
tf.concat([tf.shape(inputs)[:-1], [output_dim]], axis=0))
return act
def FProp(self, theta, inputs):
"""Apply projection to inputs.
Args:
theta: A NestedMap object containing weights' values of this layer and its
children layers.
inputs: The inputs tensor. Shaped [..., input_dims].
Returns:
Projected inputs.
"""
p = self.params
with tf.name_scope(p.name):
computation_cost.Add(
self, 'flops',
tf.reduce_prod(tf.cast(tf.shape(inputs)[:-1], tf.int64)) *
tf.cast(symbolic.ToTensor(p.input_dims * p.output_dims), tf.int64) *
2)
return py_utils.ProjectLastDim(inputs, theta.w, p.input_dims,
p.output_dims)
