diff --git a/tensorflow_federated/python/core/api/intrinsics_test.py b/tensorflow_federated/python/core/api/intrinsics_test.py
index fbca3814ad1d36bfd460121ca39606582a0102ab..5d76271910102d6ca0bdbbae6c041f46923d5433 100644
--- a/tensorflow_federated/python/core/api/intrinsics_test.py
+++ b/tensorflow_federated/python/core/api/intrinsics_test.py
@@ -126,6 +126,20 @@ class IntrinsicsTest(parameterized.TestCase):
 
     self.assert_type(foo, '(int32@CLIENTS -> {bool}@CLIENTS)')
 
+  def test_federated_map_with_client_dataset_reduce(self):
+
+    @computations.federated_computation(
+        computation_types.FederatedType(
+            computation_types.SequenceType(tf.int32), placements.CLIENTS, True))
+    def foo(ds):
+      val = intrinsics.federated_map(
+          computations.tf_computation(
+              lambda ds: ds.reduce(np.int32(0), lambda x, y: x + y)), ds)
+      self.assertIsInstance(val, value_base.Value)
+      return val
+
+    self.assert_type(foo, '(int32*@CLIENTS -> {int32}@CLIENTS)')
+
   def test_federated_map_with_client_non_all_equal_int(self):
 
     @computations.federated_computation(
diff --git a/tensorflow_federated/python/core/impl/compiler/BUILD b/tensorflow_federated/python/core/impl/compiler/BUILD
index 868a5f092a8f988749d831974300459b3afe9c59..732bedca2dc9b80c892eecdfccacc83678121def 100644
--- a/tensorflow_federated/python/core/impl/compiler/BUILD
+++ b/tensorflow_federated/python/core/impl/compiler/BUILD
@@ -1,3 +1,4 @@
+load("//tensorflow_federated/tools:build_defs.bzl", "py_cpu_gpu_test")
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
 
 package_group(
@@ -366,7 +367,7 @@ py_library(
     ],
 )
 
-py_test(
+py_cpu_gpu_test(
     name = "tensorflow_computation_factory_test",
     srcs = ["tensorflow_computation_factory_test.py"],
     python_version = "PY3",
diff --git a/tensorflow_federated/python/core/impl/compiler/tensorflow_computation_factory_test.py b/tensorflow_federated/python/core/impl/compiler/tensorflow_computation_factory_test.py
index 56af9a98c66de86b5e14884048e458fc58cb6592..a891aa67b648e6ff0ad2d3ec09506897e926af99 100644
--- a/tensorflow_federated/python/core/impl/compiler/tensorflow_computation_factory_test.py
+++ b/tensorflow_federated/python/core/impl/compiler/tensorflow_computation_factory_test.py
@@ -155,14 +155,14 @@ class CreateBinaryOperatorWithUpcastTest(parameterized.TestCase):
            computation_types.TensorType(tf.int32, shape=[1]),
            computation_types.TensorType(tf.int32),
        ]),
-       [tf.constant(1, shape=[1]), 2], 3),
+       [np.array([1]), 2], 3),
       ('add_int_different_types', tf.math.add,
        computation_types.StructType([
            computation_types.StructType([
                computation_types.TensorType(tf.int32, shape=[1])]),
            computation_types.TensorType(tf.int32),
        ]),
-       [[tf.constant(1, shape=[1])], 2],
+       [[np.array([1])], 2],
        structure.Struct([(None, 3)])),
       ('multiply_int_same_shape', tf.math.multiply,
        computation_types.StructType([
@@ -175,14 +175,14 @@ class CreateBinaryOperatorWithUpcastTest(parameterized.TestCase):
            computation_types.TensorType(tf.int32, shape=[1]),
            computation_types.TensorType(tf.int32),
        ]),
-       [tf.constant(1, shape=[1]), 2], 2),
+       [np.array([1]), 2], 2),
       ('multiply_int_different_types', tf.math.multiply,
        computation_types.StructType([
            computation_types.StructType([
                computation_types.TensorType(tf.int32, shape=[1])]),
            computation_types.TensorType(tf.int32)
        ]),
-       [[tf.constant(1, shape=[1])], 2],
+       [[np.array([1])], 2],
        structure.Struct([(None, 2)])),
       ('divide_int_same_shape', tf.math.divide,
        computation_types.StructType([
@@ -195,19 +195,24 @@ class CreateBinaryOperatorWithUpcastTest(parameterized.TestCase):
            computation_types.TensorType(tf.int32, shape=[1]),
            computation_types.TensorType(tf.int32),
        ]),
-       [tf.constant(1, shape=[1]), 2], 0.5),
+       [np.array([1]), 2], 0.5),
       ('divide_int_different_types', tf.math.divide,
        computation_types.StructType([
            computation_types.StructType([
                computation_types.TensorType(tf.int32, shape=[1])]),
            computation_types.TensorType(tf.int32),
        ]),
-       [[tf.constant(1, shape=[1])], 2],
+       [[np.array([1])], 2],
        structure.Struct([(None, 0.5)])),
   )
   # pyformat: enable
   def test_returns_computation(self, operator, type_signature, operands,
                                expected_result):
+    # TODO(b/142795960): arguments in parameterized are called before test main.
+    # `tf.constant` will error out on GPU and TPU without proper initialization.
+    # A suggested workaround is to use numpy as argument and transform to TF
+    # tensor inside the function.
+    operands = tf.nest.map_structure(tf.constant, operands)
     proto, _ = tensorflow_computation_factory.create_binary_operator_with_upcast(
         type_signature, operator)
 
@@ -323,6 +328,11 @@ class CreateComputationForPyFnTest(parameterized.TestCase):
        computation_types.TensorType(tf.int32),
        10,
        11),
+      ('dataset_reduce',
+       lambda ds: ds.reduce(np.int32(0), lambda x, y: x + y),
+       computation_types.SequenceType(tf.int32),
+       list(range(10)),
+       45),
   )
   # pyformat: enable
   def test_returns_computation(self, py_fn, type_signature, arg,