Improvement to eager linear regression benchmark

Before:
entry {
  name: "EagerLinearRegressionBenchmark.eager_train_cpu"
  iters: 2000
  wall_time: 2.45178794861
  extras {
    key: "examples_per_sec"
    value {
      double_value: 52206.7987456
    }
  }
}

After:
entry {
  name: "EagerLinearRegressionBenchmark.eager_train_cpu"
  iters: 2000
  wall_time: 1.9873790741
  extras {
    key: "examples_per_sec"
    value {
      double_value: 64406.4344182
    }
  }
}
PiperOrigin-RevId: 188068838

2 files changed, 11 insertions, 7 deletions

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index c6d16a3bc0..15f72786de 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -129,10 +129,10 @@ class Layer(checkpointable.CheckpointableBase):
     self._reuse = kwargs.get('_reuse')
     self._graph = None  # Will be set at build time.
     self._dtype = None if dtype is None else dtypes.as_dtype(dtype).name
-    call_fn_args = estimator_util.fn_args(self.call)
-    self._compute_previous_mask = ('mask' in call_fn_args or
+    self._call_fn_args = estimator_util.fn_args(self.call)
+    self._compute_previous_mask = ('mask' in self._call_fn_args or
                                    hasattr(self, 'compute_mask'))
-    self._call_has_scope_arg = 'scope' in call_fn_args
+    self._call_has_scope_arg = 'scope' in self._call_fn_args
 
     # These lists will be filled via successive calls
     # to self._add_inbound_node().
@@ -642,8 +642,9 @@ class Layer(checkpointable.CheckpointableBase):
     if (not hasattr(self, '_compute_previous_mask') or
         self._compute_previous_mask):
       previous_mask = _collect_previous_mask(inputs)
-      if ('mask' in estimator_util.fn_args(self.call) and
-          'mask' not in kwargs and
+      if not hasattr(self, '_call_fn_args'):
+        self._call_fn_args = estimator_util.fn_args(self.call)
+      if ('mask' in self._call_fn_args and 'mask' not in kwargs and
           not _is_all_none(previous_mask)):
         # The previous layer generated a mask, and mask was not explicitly pass
         # to __call__, hence we set previous_mask as the default value.
@@ -699,7 +700,9 @@ class Layer(checkpointable.CheckpointableBase):
           # TODO(agarwal): Fix the sub-classes and avoid this complexity.
           call_has_scope_arg = self._call_has_scope_arg
         except AttributeError:
-          call_has_scope_arg = 'scope' in estimator_util.fn_args(self.call)
+          self._call_fn_args = estimator_util.fn_args(self.call)
+          self._call_has_scope_arg = 'scope' in self._call_fn_args
+          call_has_scope_arg = self._call_has_scope_arg
         if call_has_scope_arg:
           kwargs['scope'] = scope
         # Check input assumptions set after layer building, e.g. input shape.
diff --git a/tensorflow/python/layers/core.py b/tensorflow/python/layers/core.py
index 6970bf9234..bdbbc59eaf 100644
--- a/tensorflow/python/layers/core.py
+++ b/tensorflow/python/layers/core.py
@@ -35,6 +35,7 @@ from tensorflow.python.layers import utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import gen_math_ops
 from tensorflow.python.ops import nn
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import standard_ops
@@ -159,7 +160,7 @@ class Dense(base.Layer):
         output_shape = shape[:-1] + [self.units]
         outputs.set_shape(output_shape)
     else:
-      outputs = standard_ops.matmul(inputs, self.kernel)
+      outputs = gen_math_ops.mat_mul(inputs, self.kernel)
     if self.use_bias:
       outputs = nn.bias_add(outputs, self.bias)
     if self.activation is not None: