Limit dtype promotion in Constmant mean (#573)

* test: Add failing test for Constant mean function dtype preservation (#523)

* fix: Preserve dtype in Constant mean function (#523)

* format: run formatter on mean tests

Author: sethaxen

gpjax/mean_functions.py

@@ -147,9 +147,9 @@ def __call__(self, x: Num[Array, "N D"]) -> Float[Array, "N O"]:
             Float[Array, "1"]: The evaluated mean function.
         """
         if isinstance(self.constant, Parameter):
-            return jnp.ones((x.shape[0], 1)) * self.constant.value
+            return jnp.ones((x.shape[0], 1), dtype=x.dtype) * self.constant.value
         else:
-            return jnp.ones((x.shape[0], 1)) * self.constant
+            return jnp.ones((x.shape[0], 1), dtype=x.dtype) * self.constant
 
 
 class Zero(Constant):

tests/test_mean_functions.py

@@ -38,7 +38,10 @@
     Constant,
     Zero,
 )
-from gpjax.parameters import Parameter
+from gpjax.parameters import (
+    Parameter,
+    Real,
+)
 
 
 def test_abstract() -> None:
@@ -323,3 +326,15 @@ def test_zero_mean_function_uses_raw_value():
     result = meanf(x)
     expected = jnp.array([[0.0], [0.0], [0.0]])
     assert jnp.allclose(result, expected)
+
+
+@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64])
+@pytest.mark.parametrize("partype", [Real, jnp.array])
+def test_constant_dtype_preservation(dtype, partype):
+    """Test that Constant mean function preserves dtype of the constant."""
+    x = jnp.arange(5, dtype=dtype).reshape(-1, 1)
+    constant = partype(jnp.array(3.0, dtype=dtype))
+    mean_fn = Constant(constant)
+    mean = mean_fn(x)
+
+    assert mean.dtype == dtype