diff --git a/python/mpact/mpactbackend.py b/python/mpact/mpactbackend.py
index 9b33892..d944376 100644
--- a/python/mpact/mpactbackend.py
+++ b/python/mpact/mpactbackend.py
@@ -359,6 +359,12 @@
         raise RuntimeError(f"Unsupported sparse layout for {a}")
 
 
+def sparse_arg(args, i):
+    if isinstance(args[i], torch.fx.node.Node):
+        return args[i].meta.get("sparsity", None)
+    return None
+
+
 def sparse_export(
     f: Callable, args: Tuple[Any, ...], kwargs: Optional[Dict[str, Any]] = None
 ) -> torch.export.ExportedProgram:
@@ -402,8 +408,22 @@
             # TODO: use upstream _opname implementation when available
             opname = node.target._schema.name.split("::")[1]
             # Zero preserving elt-wise unary op.
-            if opname in {"abs", "neg", "relu", "sin", "mul"}:
-                node.meta["sparsity"] = node.args[0].meta.get("sparsity", None)
+            if opname in {"abs", "neg", "relu", "sin"}:
+                node.meta["sparsity"] = sparse_arg(node.args, 0)
+            # Some simplistic rules for preserving sparsity. Soon
+            # to be replaced by proper FX graph propagation.
+            elif opname in {"mul"}:
+                m0 = sparse_arg(node.args, 0)
+                m1 = sparse_arg(node.args, 1)
+                if m0 is not None:
+                    node.meta["sparsity"] = m0
+                elif m1 is not None:
+                    node.meta["sparsity"] = m1
+            elif opname in {"add", "mm"}:
+                m0 = sparse_arg(node.args, 0)
+                m1 = sparse_arg(node.args, 1)
+                if m0 is not None and m1 is not None:
+                    node.meta["sparsity"] = m0
             elif opname == "_to_sparse" or opname == "to_sparse":
                 dim = len(node.meta.get("val").shape)
                 node.meta["sparsity"] = SparsityMeta(
@@ -412,13 +432,13 @@
             # TODO: Uncomment this to hack sparsity into the network.
             # elif opname == "_to_dense" or opname == "to_dense":
             #     # hack (assumes we never really want the to_dense for now)
-            #     node.meta["sparsity"] = node.args[0].meta.get("sparsity", None)
-            elif opname == "select" and node.args[0].meta.get("sparsity", None):
+            #     node.meta["sparsity"] = sparse_arg(node.args, 0)
+            elif opname == "select" and sparse_arg(node.args, 0):
                 dim = len(node.meta.get("val").shape)
                 node.meta["sparsity"] = SparsityMeta(
                     torch.sparse_coo, 0, dim, 0, None, torch.int64, torch.int64
                 )
-            elif opname == "stack" and node.args[0][0].meta.get("sparsity", None):
+            elif opname == "stack" and sparse_arg(node.args[0], 0):
                 dim = len(node.meta.get("val").shape)
                 node.meta["sparsity"] = SparsityMeta(
                     torch.sparse_coo, 0, dim - 1, 1, None, torch.int64, torch.int64
diff --git a/test/python/add.py b/test/python/add.py
new file mode 100644
index 0000000..00d4d62
--- /dev/null
+++ b/test/python/add.py
@@ -0,0 +1,89 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+import torch
+import numpy as np
+
+from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+
+from mpact.models.kernels import AddNet
+
+
+def print_sparse(res):
+    print(res[0])
+    print(res[1])
+    print(res[2])
+
+
+net = AddNet()
+
+# Construct dense and sparse matrices.
+X = torch.arange(0, 16, dtype=torch.float32).view(4, 4)
+Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
+A = torch.tensor(
+    [
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0, 2.0],
+        [0.0, 0.0, 0.0, 0.0],
+        [3.0, 0.0, 0.0, 0.0],
+    ],
+    dtype=torch.float32,
+)
+S = A.to_sparse_csr()
+
+#
+# CHECK: pytorch
+# CHECK:   tensor({{\[}}[16., 18., 20., 22.],
+# CHECK:                [24., 26., 28., 30.],
+# CHECK:                [32., 34., 36., 38.],
+# CHECK:                [40., 42., 44., 46.]{{\]}})
+# CHECK:  tensor({{\[}}[16., 18., 18., 19.],
+# CHECK:               [20., 21., 22., 25.],
+# CHECK:               [24., 25., 26., 27.],
+# CHECK:               [31., 29., 30., 31.]{{\]}})
+# CHECK:  tensor({{\[}}[ 0.,  2.,  2.,  3.],
+# CHECK:               [ 4.,  5.,  6.,  9.],
+# CHECK:               [ 8.,  9., 10., 11.],
+# CHECK:               [15., 13., 14., 15.]{{\]}})
+# CHECK:  tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
+# CHECK:         col_indices=tensor([1, 3, 0]),
+# CHECK:         values=tensor([2., 4., 6.]), size=(4, 4), nnz=3,
+# CHECK:         layout=torch.sparse_csr)
+# CHECK: mpact
+# CHECK:   {{\[}}[16. 18. 20. 22.]
+# CHECK:         [24. 26. 28. 30.]
+# CHECK:         [32. 34. 36. 38.]
+# CHECK:         [40. 42. 44. 46.]{{\]}}
+# CHECK:   {{\[}}[16. 18. 18. 19.]
+# CHECK:         [20. 21. 22. 25.]
+# CHECK:         [24. 25. 26. 27.]
+# CHECK:         [31. 29. 30. 31.]{{\]}}
+# CHECK:   {{\[}}[ 0.  2.  2.  3.]
+# CHECK:         [ 4.  5.  6.  9.]
+# CHECK:         [ 8.  9. 10. 11.]
+# CHECK:         [15. 13. 14. 15.]{{\]}}
+# CHECK:  [0 1 2 2 3]
+# CHECK:  [1 3 0]
+# CHECK:  [2. 4. 6.]
+#
+
+# Run it with PyTorch.
+print("pytorch")
+res = net(X, Y)
+print(res)
+res = net(S, Y)
+print(res)
+res = net(X, S)
+print(res)
+res = net(S, S)
+print(res)
+
+# Run it with MPACT.
+print("mpact")
+res = mpact_jit(net, X, Y)
+print(res)
+res = mpact_jit(net, S, Y)
+print(res)
+res = mpact_jit(net, X, S)
+print(res)
+res = mpact_jit(net, S, S)
+print_sparse(res)
diff --git a/test/python/mm.py b/test/python/mm.py
new file mode 100644
index 0000000..3c51c37
--- /dev/null
+++ b/test/python/mm.py
@@ -0,0 +1,89 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+import torch
+import numpy as np
+
+from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+
+from mpact.models.kernels import MMNet
+
+
+def print_sparse(res):
+    print(res[0])
+    print(res[1])
+    print(res[2])
+
+
+net = MMNet()
+
+# Construct dense and sparse matrices.
+X = torch.arange(0, 16, dtype=torch.float32).view(4, 4)
+Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
+A = torch.tensor(
+    [
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0, 2.0],
+        [0.0, 0.0, 0.0, 0.0],
+        [3.0, 0.0, 0.0, 0.0],
+    ],
+    dtype=torch.float32,
+)
+S = A.to_sparse_csr()
+
+#
+# CHECK: pytorch
+# CHECK:   tensor({{\[}}[ 152.,  158.,  164.,  170.],
+# CHECK:                [ 504.,  526.,  548.,  570.],
+# CHECK:                [ 856.,  894.,  932.,  970.],
+# CHECK:                [1208., 1262., 1316., 1370.]{{\]}})
+# CHECK:  tensor({{\[}}[20., 21., 22., 23.],
+# CHECK:               [56., 58., 60., 62.],
+# CHECK:               [ 0.,  0.,  0.,  0.],
+# CHECK:               [48., 51., 54., 57.]{{\]}})
+# CHECK:  tensor({{\[}}[ 9.,  0.,  0.,  2.],
+# CHECK:               [21.,  4.,  0., 10.],
+# CHECK:               [33.,  8.,  0., 18.],
+# CHECK:               [45., 12.,  0., 26.]{{\]}})
+# CHECK:  tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
+# CHECK:         col_indices=tensor([3, 0, 1]),
+# CHECK:         values=tensor([2., 6., 3.]), size=(4, 4), nnz=3,
+# CHECK:         layout=torch.sparse_csr)
+# CHECK: mpact
+# CHECK:   {{\[}}[ 152.  158.  164.  170.]
+# CHECK:         [ 504.  526.  548.  570.]
+# CHECK:         [ 856.  894.  932.  970.]
+# CHECK:         [1208. 1262. 1316. 1370.]{{\]}}
+# CHECK:   {{\[}}[20. 21. 22. 23.]
+# CHECK:         [56. 58. 60. 62.]
+# CHECK:         [ 0.  0.  0.  0.]
+# CHECK:         [48. 51. 54. 57.]{{\]}}
+# CHECK:   {{\[}}[ 9.  0.  0.  2.]
+# CHECK:         [21.  4.  0. 10.]
+# CHECK:         [33.  8.  0. 18.]
+# CHECK:         [45. 12.  0. 26.]{{\]}}
+# CHECK:  [0 1 2 2 3]
+# CHECK:  [3 0 1]
+# CHECK:  [2. 6. 3.]
+#
+
+# Run it with PyTorch.
+print("pytorch")
+res = net(X, Y)
+print(res)
+res = net(S, Y)
+print(res)
+res = net(X, S)
+print(res)
+res = net(S, S)
+print(res)
+
+# Run it with MPACT.
+print("mpact")
+res = mpact_jit(net, X, Y)
+print(res)
+res = mpact_jit(net, S, Y)
+print(res)
+res = mpact_jit(net, X, S)
+print(res)
+res = mpact_jit(net, S, S)
+print_sparse(res)
diff --git a/test/python/mul.py b/test/python/mul.py
new file mode 100644
index 0000000..fd8692f
--- /dev/null
+++ b/test/python/mul.py
@@ -0,0 +1,87 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+import torch
+import numpy as np
+
+from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+
+from mpact.models.kernels import MulNet
+
+
+def print_sparse(res):
+    print(res[0])
+    print(res[1])
+    print(res[2])
+
+
+net = MulNet()
+
+# Construct dense and sparse matrices.
+X = torch.arange(0, 16, dtype=torch.float32).view(4, 4)
+Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
+A = torch.tensor(
+    [
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0, 2.0],
+        [0.0, 0.0, 0.0, 0.0],
+        [3.0, 0.0, 0.0, 0.0],
+    ],
+    dtype=torch.float32,
+)
+S = A.to_sparse_csr()
+
+#
+# CHECK: pytorch
+# CHECK: tensor({{\[}}[  0.,  17.,  36.,  57.],
+# CHECK:              [ 80., 105., 132., 161.],
+# CHECK:              [192., 225., 260., 297.],
+# CHECK:              [336., 377., 420., 465.]{{\]}})
+# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
+# CHECK:        col_indices=tensor([1, 3, 0]),
+# CHECK:        values=tensor([17., 46., 84.]), size=(4, 4), nnz=3,
+# CHECK:        layout=torch.sparse_csr)
+# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
+# CHECK:        col_indices=tensor([1, 3, 0]),
+# CHECK:        values=tensor([ 1., 14., 36.]), size=(4, 4), nnz=3,
+# CHECK:        layout=torch.sparse_csr)
+# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
+# CHECK:        col_indices=tensor([1, 3, 0]),
+# CHECK:        values=tensor([1., 4., 9.]), size=(4, 4), nnz=3,
+# CHECK:        layout=torch.sparse_csr)
+# CHECK: mpact
+# CHECK:   {{\[}}[  0.  17.  36.  57.]
+# CHECK:         [ 80. 105. 132. 161.]
+# CHECK:         [192. 225. 260. 297.]
+# CHECK:         [336. 377. 420. 465.]{{\]}}
+# CHECK:  [0 1 2 2 3]
+# CHECK:  [1 3 0]
+# CHECK:  [17. 46. 84.]
+# CHECK:  [0 1 2 2 3]
+# CHECK:  [1 3 0]
+# CHECK:  [ 1. 14. 36.]
+# CHECK:  [0 1 2 2 3]
+# CHECK:  [1 3 0]
+# CHECK:  [1. 4. 9.]
+#
+
+# Run it with PyTorch.
+print("pytorch")
+res = net(X, Y)
+print(res)
+res = net(S, Y)
+print(res)
+res = net(X, S)
+print(res)
+res = net(S, S)
+print(res)
+
+# Run it with MPACT.
+print("mpact")
+res = mpact_jit(net, X, Y)
+print(res)
+res = mpact_jit(net, S, Y)
+print_sparse(res)
+res = mpact_jit(net, X, S)
+print_sparse(res)
+res = mpact_jit(net, S, S)
+print_sparse(res)
diff --git a/test/python/sddmm.py b/test/python/sddmm.py
index 8793120..a0890c7 100644
--- a/test/python/sddmm.py
+++ b/test/python/sddmm.py
@@ -7,6 +7,14 @@
 
 from mpact.models.kernels import MMNet, SDDMMNet
 
+
+def print_sparse(res):
+    print(res[0])
+    print(res[1])
+    print(res[2])
+    print(res[3])
+
+
 mmnet = MMNet()
 sddmmnet = SDDMMNet()
 
@@ -54,7 +62,4 @@
 dense = mpact_jit(mmnet, A, B)
 print(dense)
 res = mpact_jit(sddmmnet, S, A, B)
-print(res[0])
-print(res[1])
-print(res[2])
-print(res[3])
+print_sparse(res)