Rename b to x, matching BLAS docs

Author: jessegrabowski

pytensor/tensor/slinalg.py

@@ -1678,20 +1678,20 @@ def __init__(self, lower_diags, upper_diags):
         self.lower_diags = lower_diags
         self.upper_diags = upper_diags
 
-    def make_node(self, A, b):
+    def make_node(self, A, x):
         A = as_tensor_variable(A)
-        B = as_tensor_variable(b)
+        x = as_tensor_variable(x)
 
-        out_dtype = pytensor.scalar.upcast(A.dtype, B.dtype)
-        output = b.type.clone(dtype=out_dtype)()
+        out_dtype = pytensor.scalar.upcast(A.dtype, x.dtype)
+        output = x.type.clone(dtype=out_dtype)()
 
-        return pytensor.graph.basic.Apply(self, [A, B], [output])
+        return pytensor.graph.basic.Apply(self, [A, x], [output])
 
     def infer_shape(self, fgraph, nodes, shapes):
         return [shapes[0][:-1]]
 
     def perform(self, node, inputs, outputs_storage):
-        A, b = inputs
+        A, x = inputs
         m, n = A.shape
 
         kl = self.lower_diags
@@ -1703,10 +1703,10 @@ def perform(self, node, inputs, outputs_storage):
             A_banded[i, slice(k, None) if k >= 0 else slice(None, n + k)] = diag(A, k=k)
 
         fn = scipy_linalg.get_blas_funcs("gbmv", dtype=A.dtype)
-        outputs_storage[0][0] = fn(m=m, n=n, kl=kl, ku=ku, alpha=1, a=A_banded, x=b)
+        outputs_storage[0][0] = fn(m=m, n=n, kl=kl, ku=ku, alpha=1, a=A_banded, x=x)
 
 
-def banded_dot(A: TensorLike, b: TensorLike, lower_diags: int, upper_diags: int):
+def banded_dot(A: TensorLike, x: TensorLike, lower_diags: int, upper_diags: int):
     """
     Specialized matrix-vector multiplication for cases when A is a banded matrix
 
@@ -1719,7 +1719,7 @@ def banded_dot(A: TensorLike, b: TensorLike, lower_diags: int, upper_diags: int)
     ----------
     A: Tensorlike
         Matrix to perform banded dot on.
-    b: Tensorlike
+    x: Tensorlike
         Vector to perform banded dot on.
     lower_diags: int
         Number of nonzero lower diagonals of A
@@ -1731,7 +1731,7 @@ def banded_dot(A: TensorLike, b: TensorLike, lower_diags: int, upper_diags: int)
     out: Tensor
         The matrix multiplication result
     """
-    return Blockwise(BandedDot(lower_diags, upper_diags))(A, b)
+    return Blockwise(BandedDot(lower_diags, upper_diags))(A, x)
 
 
 __all__ = [

tests/tensor/test_slinalg.py

@@ -1074,19 +1074,19 @@ def test_banded_dot(A_shape, kl, ku):
     rng = np.random.default_rng()
 
     A_val = _make_banded_A(rng.normal(size=A_shape), kl=kl, ku=ku).astype(config.floatX)
-    b_val = rng.normal(size=(A_shape[-1],)).astype(config.floatX)
+    x_val = rng.normal(size=(A_shape[-1],)).astype(config.floatX)
 
     A = pt.tensor("A", shape=A_val.shape, dtype=A_val.dtype)
-    b = pt.tensor("b", shape=b_val.shape, dtype=b_val.dtype)
-    res = banded_dot(A, b, kl, ku)
-    res_2 = A @ b
+    x = pt.tensor("x", shape=x_val.shape, dtype=x_val.dtype)
+    res = banded_dot(A, x, kl, ku)
+    res_2 = A @ x
 
-    fn = function([A, b], [res, res_2], trust_input=True)
+    fn = function([A, x], [res, res_2], trust_input=True)
     assert any(isinstance(node.op, BandedDot) for node in fn.maker.fgraph.apply_nodes)
 
-    x_val, x2_val = fn(A_val, b_val)
+    out_val, out_2_val = fn(A_val, x_val)
 
-    np.testing.assert_allclose(x_val, x2_val)
+    np.testing.assert_allclose(out_val, out_2_val)
 
 
 @pytest.mark.parametrize("op", ["dot", "banded_dot"], ids=str)
@@ -1099,17 +1099,17 @@ def test_banded_dot_perf(op, A_shape, benchmark):
     rng = np.random.default_rng()
 
     A_val = _make_banded_A(rng.normal(size=A_shape), kl=1, ku=1).astype(config.floatX)
-    b_val = rng.normal(size=(A_shape[-1],)).astype(config.floatX)
+    x_val = rng.normal(size=(A_shape[-1],)).astype(config.floatX)
 
     A = pt.tensor("A", shape=A_val.shape, dtype=A_val.dtype)
-    b = pt.tensor("b", shape=b_val.shape, dtype=A_val.dtype)
+    x = pt.tensor("x", shape=x_val.shape, dtype=x_val.dtype)
 
     if op == "dot":
         f = pt.dot
     elif op == "banded_dot":
         f = functools.partial(banded_dot, lower_diags=1, upper_diags=1)
 
-    res = f(A, b)
-    fn = function([A, b], res, trust_input=True)
+    res = f(A, x)
+    fn = function([A, x], res, trust_input=True)
 
-    benchmark(fn, A_val, b_val)
+    benchmark(fn, A_val, x_val)