diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index c8143c0..e306dc9 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1011,6 +1011,7 @@
         "//tensorflow/python:framework",
         "//tensorflow/python:math_ops",
         "//tensorflow/python:platform_test",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
 
diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py
index 200851e..d81a06a 100644
--- a/tensorflow/compiler/tests/ternary_ops_test.py
+++ b/tensorflow/compiler/tests/ternary_ops_test.py
@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
@@ -28,7 +29,7 @@
 from tensorflow.python.platform import googletest
 
 
-class TernaryOpsTest(xla_test.XLATestCase):
+class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
 
   def _testTernary(self, op, a, b, c, expected):
     with self.session() as session:
@@ -39,20 +40,24 @@
         output = op(pa, pb, pc)
       result = session.run(output, {pa: a, pb: b, pc: c})
       self.assertAllClose(result, expected, rtol=1e-3)
+      return result
 
-  def testLinspace(self):
-    self._testTernary(
+  @parameterized.parameters(
+      {'start': 1, 'end': 2, 'num': 1},
+      {'start': 1, 'end': 4, 'num': 3},
+      {'start': 0, 'end': 41, 'num': 42})
+  def testLinspace(self, start, end, num):
+    expected = np.linspace(start, end, num, dtype=np.float32)
+    result = self._testTernary(
         math_ops.linspace,
-        np.float32(1),
-        np.float32(2),
-        np.int32(1),
-        expected=np.array([1], dtype=np.float32))
-    self._testTernary(
-        math_ops.linspace,
-        np.float32(1),
-        np.float32(4),
-        np.int32(3),
-        expected=np.array([1, 2.5, 4], dtype=np.float32))
+        np.float32(start),
+        np.float32(end),
+        np.int32(num),
+        expected)
+    # According to linspace spec, start has to be the first element and end has
+    # to be last element.
+    self.assertEqual(result[-1], expected[-1])
+    self.assertEqual(result[0], expected[0])
 
   def testRange(self):
     self._testTernary(
diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
index 1cbb142..2eef554 100644
--- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
@@ -18,6 +18,7 @@
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/client/lib/constants.h"
 #include "tensorflow/compiler/xla/client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal.h"
 #include "tensorflow/compiler/xla/primitive_util.h"
@@ -153,67 +154,27 @@
                 errors::InvalidArgument("num must be a scalar, not shape ",
                                         num_in_shape.DebugString()));
 
-    DataType type = ctx->input_type(0);
-
     int64 num;
     OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar("num", &num));
     OP_REQUIRES(ctx, num > 0,
                 errors::InvalidArgument("Requires num > 0: ", num));
-    Tensor out_constant(type, TensorShape({num}));
-
-    xla::Literal start_literal;
-    OP_REQUIRES_OK(ctx, ctx->ConstantInput("start", &start_literal));
-    xla::Literal stop_literal;
-    OP_REQUIRES_OK(ctx, ctx->ConstantInput("stop", &stop_literal));
-
-    switch (type) {
-      case DT_FLOAT: {
-        float start = start_literal.GetFirstElement<float>();
-        float stop = stop_literal.GetFirstElement<float>();
-        auto flat = out_constant.flat<float>();
-        if (num == 1) {
-          flat(0) = start;
-        } else {
-          const float step = (stop - start) / (num - 1);
-          for (int64 i = 0; i < num - 1; ++i) {
-            flat(i) = start + step * i;
-          }
-          // The last value in the sequence must be equal to stop.
-          flat(num - 1) = stop;
-        }
-        break;
-      }
-      case DT_DOUBLE: {
-        double start = start_literal.GetFirstElement<double>();
-        double stop = stop_literal.GetFirstElement<double>();
-        auto flat = out_constant.flat<double>();
-        if (num == 1) {
-          flat(0) = start;
-        } else {
-          const double step = (stop - start) / (num - 1);
-          for (int64 i = 0; i < num - 1; ++i) {
-            flat(i) = start + step * i;
-          }
-          // The last value in the sequence must be equal to stop.
-          flat(num - 1) = stop;
-        }
-        break;
-      }
-
-      default:
-        ctx->SetStatus(errors::InvalidArgument("Invalid argument type ",
-                                               DataTypeString(type)));
-        return;
+    xla::XlaOp start = ctx->Input("start");
+    xla::XlaOp stop = ctx->Input("stop");
+    xla::XlaOp iota = xla::Iota(ctx->builder(), ctx->output_xla_type(0), num) /
+                      xla::ScalarLike(start, (num > 1 ? num - 1 : num));
+    xla::XlaOp result = iota * stop - (iota * start - start);
+    if (num > 1) {
+      // According to linspace spec, start has to be the first element and end
+      // has to be last element.
+      xla::XlaOp mask = xla::Iota(ctx->builder(), xla::S64, num);
+      xla::XlaOp eq = xla::Eq(mask, xla::ScalarLike(mask, num - 1));
+      result = xla::Select(eq, stop, result);
     }
-    ctx->SetConstantOutput(0, out_constant);
+    ctx->SetOutput(0, result);
   }
 };
 
-REGISTER_XLA_OP(Name("LinSpace")
-                    .CompileTimeConstantInput("start")
-                    .CompileTimeConstantInput("stop")
-                    .CompileTimeConstantInput("num"),
-                LinSpaceOp);
+REGISTER_XLA_OP(Name("LinSpace").CompileTimeConstantInput("num"), LinSpaceOp);
 
 }  // namespace
 }  // namespace tensorflow