Move benchmarks tests into separate folder and cleanup redundancies

Author: ricardoV94

tests/benchmarks/__init__.py

@@ -0,0 +1,127 @@
+import numpy as np
+import pytest
+
+from pytensor import In, function
+from pytensor.tensor import dot, empty, matrix, outer, scalar, tensor, vector
+from pytensor.tensor.blas_c import CGemv
+
+
+@pytest.mark.parametrize("dtype", ("float64", "float32", "mixed"))
+def test_mat_vec_dot_benchmark_numba(dtype, benchmark):
+    A = tensor("A", shape=(512, 512), dtype="float64" if dtype == "mixed" else dtype)
+    x = tensor("x", shape=(512,), dtype="float32" if dtype == "mixed" else dtype)
+    out = dot(A, x)
+
+    fn = function([A, x], out, mode="NUMBA", trust_input=True)
+
+    rng = np.random.default_rng(948)
+    A_test = rng.standard_normal(size=A.type.shape).astype(A.type.dtype)
+    x_test = rng.standard_normal(size=x.type.shape).astype(x.type.dtype)
+    np.testing.assert_allclose(fn(A_test, x_test), np.dot(A_test, x_test), atol=1e-4)
+    benchmark(fn, A_test, x_test)
+
+
+def _test_ger_benchmark(mode, n, inplace, benchmark):
+    alpha = scalar("alpha")
+    x = vector("x")
+    y = vector("y")
+    A = matrix("A")
+
+    out = alpha * outer(x, y) + A
+
+    fn = function(
+        [alpha, x, y, In(A, mutable=inplace)],
+        out,
+        mode=mode,
+        trust_input=True,
+    )
+
+    rng = np.random.default_rng([2274, n])
+    alpha_test = rng.normal(size=())
+    x_test = rng.normal(size=(n,))
+    y_test = rng.normal(size=(n,))
+    A_test = rng.normal(size=(n, n))
+
+    benchmark(fn, alpha_test, x_test, y_test, A_test)
+
+
+@pytest.mark.parametrize("inplace", (True, False), ids=["inplace", "no_inplace"])
+@pytest.mark.parametrize("n", [2**7, 2**9, 2**13])
+def test_ger_benchmark_c(n, inplace, benchmark):
+    _test_ger_benchmark("CVM", n, inplace, benchmark)
+
+
+@pytest.mark.parametrize("inplace", (True, False), ids=["inplace", "no_inplace"])
+@pytest.mark.parametrize("n", [2**7, 2**9, 2**13])
+def test_ger_benchmark_numba(n, inplace, benchmark):
+    _test_ger_benchmark("NUMBA", n, inplace, benchmark)
+
+
+def test_cgemv_vector_dot_benchmark(benchmark):
+    n = 400_000
+    a = vector("A", shape=(n,))
+    b = vector("x", shape=(n,))
+
+    out = CGemv(inplace=True)(
+        empty((1,)),
+        1.0,
+        a[None],
+        b,
+        0.0,
+    )
+    fn = function([a, b], out, accept_inplace=True, mode="CVM", trust_input=True)
+
+    rng = np.random.default_rng(430)
+    test_a = rng.normal(size=n)
+    test_b = rng.normal(size=n)
+
+    np.testing.assert_allclose(
+        fn(test_a, test_b),
+        np.dot(test_a, test_b),
+    )
+
+    benchmark(fn, test_a, test_b)
+
+
+@pytest.mark.parametrize(
+    "neg_stride1", (True, False), ids=["neg_stride1", "pos_stride1"]
+)
+@pytest.mark.parametrize(
+    "neg_stride0", (True, False), ids=["neg_stride0", "pos_stride0"]
+)
+@pytest.mark.parametrize("F_layout", (True, False), ids=["F_layout", "C_layout"])
+def test_cgemv_negative_strides_benchmark(
+    neg_stride0, neg_stride1, F_layout, benchmark
+):
+    A = matrix("A", shape=(512, 512))
+    x = vector("x", shape=(A.type.shape[-1],))
+    y = vector("y", shape=(A.type.shape[0],))
+
+    out = CGemv(inplace=False)(
+        y,
+        1.0,
+        A,
+        x,
+        1.0,
+    )
+    fn = function([A, x, y], out, trust_input=True, mode="CVM")
+
+    rng = np.random.default_rng(430)
+    test_A = rng.normal(size=A.type.shape)
+    test_x = rng.normal(size=x.type.shape)
+    test_y = rng.normal(size=y.type.shape)
+
+    if F_layout:
+        test_A = test_A.T
+    if neg_stride0:
+        test_A = test_A[::-1]
+    if neg_stride1:
+        test_A = test_A[:, ::-1]
+    assert (test_A.strides[0] < 0) == neg_stride0
+    assert (test_A.strides[1] < 0) == neg_stride1
+
+    # Check result is correct by using a copy of A with positive strides
+    res = fn(test_A, test_x, test_y)
+    np.testing.assert_allclose(res, fn(test_A.copy(), test_x, test_y))
+
+    benchmark(fn, test_A, test_x, test_y)

tests/benchmarks/test_blas.py

@@ -0,0 +1,93 @@
+import numpy as np
+import pytest
+
+from pytensor import function
+from pytensor.tensor import diagonal, grad, log, tensor
+from pytensor.tensor.linalg import cholesky, solve_triangular
+
+
+def _test_blockwise_cholesky_benchmark(mode, benchmark):
+    from pytensor.tensor.blockwise import Blockwise
+
+    x = tensor(shape=(5, 3, 3))
+    out = cholesky(x)
+    assert isinstance(out.owner.op, Blockwise)
+
+    fn = function([x], out, mode=mode, trust_input=True)
+    x_test = np.eye(3) * np.arange(1, 6)[:, None, None]
+    fn(x_test)  # JIT compile
+    benchmark(fn, x_test)
+
+
+def test_blockwise_cholesky_benchmark_c(benchmark):
+    _test_blockwise_cholesky_benchmark("CVM", benchmark)
+
+
+def test_blockwise_cholesky_benchmark_numba(benchmark):
+    _test_blockwise_cholesky_benchmark("NUMBA", benchmark)
+
+
+def _test_batched_mvnormal_logp_and_dlogp(
+    mode, mu_batch_shape, cov_batch_shape, benchmark
+):
+    rng = np.random.default_rng(sum(map(ord, "batched_mvnormal")))
+
+    value_batch_shape = mu_batch_shape
+    if len(cov_batch_shape) > len(mu_batch_shape):
+        value_batch_shape = cov_batch_shape
+
+    value = tensor("value", shape=(*value_batch_shape, 10))
+    mu = tensor("mu", shape=(*mu_batch_shape, 10))
+    cov = tensor("cov", shape=(*cov_batch_shape, 10, 10))
+
+    test_values = [
+        rng.normal(size=value.type.shape),
+        rng.normal(size=mu.type.shape),
+        np.eye(cov.type.shape[-1]) * np.abs(rng.normal(size=cov.type.shape)),
+    ]
+
+    chol_cov = cholesky(cov, lower=True, on_error="raise")
+    delta_trans = solve_triangular(chol_cov, value - mu, b_ndim=1)
+    quaddist = (delta_trans**2).sum(axis=-1)
+    diag = diagonal(chol_cov, axis1=-2, axis2=-1)
+    logdet = log(diag).sum(axis=-1)
+    k = value.shape[-1]
+    norm = -0.5 * k * (np.log(2 * np.pi))
+
+    logp = norm - 0.5 * quaddist - logdet
+    dlogp = grad(logp.sum(), wrt=[value, mu, cov])
+
+    fn = function([value, mu, cov], [logp, *dlogp], mode=mode, trust_input=True)
+    benchmark(fn, *test_values)
+
+
+@pytest.mark.parametrize(
+    "mu_batch_shape", [(), (1000,), (4, 1000)], ids=lambda arg: f"mu:{arg}"
+)
+@pytest.mark.parametrize(
+    "cov_batch_shape", [(), (1000,), (4, 1000)], ids=lambda arg: f"cov:{arg}"
+)
+def test_batched_mvnormal_logp_and_dlogp_c(mu_batch_shape, cov_batch_shape, benchmark):
+    _test_batched_mvnormal_logp_and_dlogp(
+        mode="CVM",
+        mu_batch_shape=mu_batch_shape,
+        cov_batch_shape=cov_batch_shape,
+        benchmark=benchmark,
+    )
+
+
+@pytest.mark.parametrize(
+    "mu_batch_shape", [(), (1000,), (4, 1000)], ids=lambda arg: f"mu:{arg}"
+)
+@pytest.mark.parametrize(
+    "cov_batch_shape", [(), (1000,), (4, 1000)], ids=lambda arg: f"cov:{arg}"
+)
+def test_batched_mvnormal_logp_and_dlogp_numba(
+    mu_batch_shape, cov_batch_shape, benchmark
+):
+    _test_batched_mvnormal_logp_and_dlogp(
+        mode="NUMBA",
+        mu_batch_shape=mu_batch_shape,
+        cov_batch_shape=cov_batch_shape,
+        benchmark=benchmark,
+    )

tests/benchmarks/test_blockwise.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pytest
+
+from pytensor import function, shared
+
+
+def _test_careduce_benchmark(axis, c_contiguous, mode, benchmark):
+    N = 256
+    x_test = np.random.uniform(size=(N, N, N))
+    transpose_axis = (0, 1, 2) if c_contiguous else (2, 0, 1)
+
+    x = shared(x_test, name="x", shape=x_test.shape)
+    out = x.transpose(transpose_axis).sum(axis=axis)
+    fn = function([], out, mode=mode)
+
+    np.testing.assert_allclose(
+        fn(),
+        x_test.transpose(transpose_axis).sum(axis=axis),
+    )
+    benchmark(fn)
+
+
+@pytest.mark.parametrize(
+    "axis",
+    (0, 1, 2, (0, 1), (0, 2), (1, 2), None),
+    ids=lambda x: f"axis={x}",
+)
+@pytest.mark.parametrize(
+    "c_contiguous",
+    (True, False),
+    ids=lambda x: f"c_contiguous={x}",
+)
+def test_careduce_benchmark_c(axis, c_contiguous, benchmark):
+    _test_careduce_benchmark(
+        axis=axis, c_contiguous=c_contiguous, mode="CVM", benchmark=benchmark
+    )
+
+
+@pytest.mark.parametrize(
+    "axis",
+    (0, 1, 2, (0, 1), (0, 2), (1, 2), None),
+    ids=lambda x: f"axis={x}",
+)
+@pytest.mark.parametrize(
+    "c_contiguous",
+    (True, False),
+    ids=lambda x: f"c_contiguous={x}",
+)
+def test_careduce_benchmark_numba(axis, c_contiguous, benchmark):
+    _test_careduce_benchmark(
+        axis=axis, c_contiguous=c_contiguous, mode="NUMBA", benchmark=benchmark
+    )