Add float32 test support

Author: jessegrabowski

tests/tensor/test_einsum.py

@@ -14,6 +14,10 @@
 from pytensor.tensor.shape import Reshape
 
 
+floatX = pytensor.config.floatX
+ATOL = RTOL = 1e-8 if floatX == "float64" else 1e-4
+
+
 def assert_no_blockwise_in_graph(fgraph: FunctionGraph, core_op=None) -> None:
     for node in fgraph.apply_nodes:
         if isinstance(node.op, Blockwise):
@@ -79,11 +83,10 @@ def test_general_dot():
     np_batched_tensordot = np.vectorize(
         partial(np.tensordot, axes=tensordot_axes), signature=signature
     )
-    x_test = rng.normal(size=x.type.shape)
-    y_test = rng.normal(size=y.type.shape)
+    x_test = rng.normal(size=x.type.shape).astype(floatX)
+    y_test = rng.normal(size=y.type.shape).astype(floatX)
     np.testing.assert_allclose(
-        fn(x_test, y_test),
-        np_batched_tensordot(x_test, y_test),
+        fn(x_test, y_test), np_batched_tensordot(x_test, y_test), atol=ATOL, rtol=RTOL
     )
 
 
@@ -130,7 +133,7 @@ def test_einsum_signatures(static_shape_known, signature):
     assert out.owner.op.optimize == static_shape_known or len(operands) <= 2
 
     rng = np.random.default_rng(37)
-    test_values = [rng.normal(size=shape) for shape in shapes]
+    test_values = [rng.normal(size=shape).astype(floatX) for shape in shapes]
     np_out = np.einsum(signature, *test_values)
 
     fn = function(operands, out)
@@ -139,7 +142,7 @@ def test_einsum_signatures(static_shape_known, signature):
     # print(); fn.dprint(print_type=True)
 
     assert_no_blockwise_in_graph(fn.maker.fgraph)
-    np.testing.assert_allclose(pt_out, np_out)
+    np.testing.assert_allclose(pt_out, np_out, atol=ATOL, rtol=RTOL)
 
 
 def test_batch_dim():
@@ -165,40 +168,49 @@ def test_einsum_conv():
     conv_signature = "bchwkt,fckt->bfhw"
     windowed_input = rng.random(
         size=(batch_size, channels, height, width, kernel_size, kernel_size)
+    ).astype(floatX)
+    weights = rng.random(size=(num_filters, channels, kernel_size, kernel_size)).astype(
+        floatX
     )
-    weights = rng.random(size=(num_filters, channels, kernel_size, kernel_size))
     result = einsum(conv_signature, windowed_input, weights).eval()
 
     assert result.shape == (32, 15, 8, 8)
     np.testing.assert_allclose(
         result,
         np.einsum("bchwkt,fckt->bfhw", windowed_input, weights),
+        atol=ATOL,
+        rtol=RTOL,
     )
 
 
 def test_ellipsis():
     rng = np.random.default_rng(159)
     x = pt.tensor("x", shape=(3, 5, 7, 11))
     y = pt.tensor("y", shape=(3, 5, 11, 13))
-    x_test = rng.normal(size=x.type.shape)
-    y_test = rng.normal(size=y.type.shape)
+    x_test = rng.normal(size=x.type.shape).astype(floatX)
+    y_test = rng.normal(size=y.type.shape).astype(floatX)
     expected_out = np.matmul(x_test, y_test)
 
     with pytest.raises(ValueError):
         pt.einsum("mp,pn->mn", x, y)
 
     out = pt.einsum("...mp,...pn->...mn", x, y)
-    np.testing.assert_allclose(out.eval({x: x_test, y: y_test}), expected_out)
+    np.testing.assert_allclose(
+        out.eval({x: x_test, y: y_test}), expected_out, atol=ATOL, rtol=RTOL
+    )
 
     # Put batch axes in the middle
     new_x = pt.moveaxis(x, -2, 0)
     new_y = pt.moveaxis(y, -2, 0)
     out = pt.einsum("m...p,p...n->m...n", new_x, new_y)
     np.testing.assert_allclose(
-        out.eval({x: x_test, y: y_test}), expected_out.transpose(-2, 0, 1, -1)
+        out.eval({x: x_test, y: y_test}),
+        expected_out.transpose(-2, 0, 1, -1),
+        atol=ATOL,
+        rtol=RTOL,
     )
 
     out = pt.einsum("m...p,p...n->mn", new_x, new_y)
     np.testing.assert_allclose(
-        out.eval({x: x_test, y: y_test}), expected_out.sum((0, 1))
+        out.eval({x: x_test, y: y_test}), expected_out.sum((0, 1)), atol=ATOL, rtol=RTOL
     )