output_arg { name: 'b' type_attr: 'T' }
attr { name: 'T' type: 'type' }
}
op {
name: 'OpWithDefaultAttr'
output_arg { name: 'a' type: DT_INT32 }
attr { name: 'default_float' type: 'float' default_value { f: 123.0 } }
}
op {
name: 'OpWithFutureDefaultAttr'
}
""", _op_list)
op_def_registry.register_op_list(_op_list)
# NOTE(mrry): Dummy shape registrations for ops used in the tests.
for op_def in _op_list.op:
tf.RegisterShape(op_def.name)(None)
class ImportGraphDefTest(tf.test.TestCase):
def _MakeGraphDef(self,
text,
producer=tf.GRAPH_DEF_VERSION,
min_consumer=tf.GRAPH_DEF_VERSION_MIN_CONSUMER):
text = "versions: { producer: %d min_consumer: %d };\n%s" % (
producer, min_consumer, text)
ret = tf.GraphDef()
text_format.Merge(text, ret)
return ret
def testBasic(self):
with tf.Graph().as_default():
from tensorflow.contrib.framework.python.ops import add_arg_scope
from tensorflow.contrib.framework.python.ops import variables
from tensorflow.contrib.layers.python.layers import utils
from tensorflow.python.framework import ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import variable_scope
import tensorflow as tf
cMod = tf.load_op_library('layer_norm_fused_op.so')
# disabled these if using newer version of Tensorflow. (You can keep this
# if no error raised)
tf.RegisterShape("LayerNormCustom")(common_shapes.call_cpp_shape_fn)
tf.RegisterShape("LayerNormBiasAddCustom")(common_shapes.call_cpp_shape_fn)
tf.RegisterShape("LayerNormFusedCustom")(common_shapes.call_cpp_shape_fn)
@ops.RegisterGradient("LayerNormCustom")
def _LayerNormCustomGrad(op, grad):
return [
cMod.layer_norm_backprop_custom(op.inputs[0], grad,
op.get_attr("epsilon"))
]
@ops.RegisterGradient("LayerNormBiasAddCustom")
def _LayerNormBiasAddCustomGrad(op, grad):
in_back, beta_back = cMod.layer_norm_bias_add_backprop_custom(
