diff --git a/matadi/_templates.py b/matadi/_templates.py
index fe16b4d..fb24c0b 100644
--- a/matadi/_templates.py
+++ b/matadi/_templates.py
@@ -188,9 +188,9 @@ def hessian(self, x, **kwargs):
 
 
 class MaterialTensorGeneral(MaterialTensor):
-    def __init__(self, fun, nstatevars=1, x=None, triu=True, **kwargs):
+    def __init__(self, fun, statevars_shape=(1, 1), x=None, triu=True, **kwargs):
         """A (first Piola-Kirchhoff stress) tensor-based material definition with
-        ``n`` state variables."""
+        state variables of a given shape."""
 
         if x is None:
             x = [Variable("F", 3, 3)]
@@ -202,6 +202,6 @@ def __init__(self, fun, nstatevars=1, x=None, triu=True, **kwargs):
             pass
 
         # add state variables
-        x.append(Variable("z", nstatevars, 1))
+        x.append(Variable("z", *statevars_shape))
 
         super().__init__(x=x, fun=fun, triu=triu, statevars=1, kwargs=kwargs)
diff --git a/matadi/math.py b/matadi/math.py
index 421c554..480347e 100644
--- a/matadi/math.py
+++ b/matadi/math.py
@@ -74,6 +74,16 @@
 zeros = SX.zeros
 
 
+def zeros_like(T):
+    
+    return zeros(T.shape)
+
+
+def ones_like(T):
+    
+    return ones(T.shape)
+
+
 def invariants(T):
 
     I1 = trace(T)
diff --git a/matadi/models/_templates.py b/matadi/models/_templates.py
index 6273fb0..9ef0c2d 100644
--- a/matadi/models/_templates.py
+++ b/matadi/models/_templates.py
@@ -28,7 +28,7 @@ def fun(x, C10, r, m, beta, bulk):
             # for a pseudo-elastic material formulation
             return eta * gradient(W, F) + gradient(U, F), Wmax
 
-        super().__init__(fun=fun, nstatevars=1, C10=C10, r=r, m=m, beta=beta, bulk=bulk)
+        super().__init__(fun=fun, statevars_shape=(1, 1), C10=C10, r=r, m=m, beta=beta, bulk=bulk)
 
 
 class Morph(MaterialTensorGeneral):
@@ -59,7 +59,7 @@ def fun(x, p1, p2, p3, p4, p5, p6, p7, p8, bulk):
 
         super().__init__(
             fun=fun,
-            nstatevars=13,
+            statevars_shape=(13, 1),
             p1=p1,
             p2=p2,
             p3=p3,
diff --git a/matadi/models/microsphere/affine/_models.py b/matadi/models/microsphere/affine/_models.py
index 1d8aab5..727a77d 100644
--- a/matadi/models/microsphere/affine/_models.py
+++ b/matadi/models/microsphere/affine/_models.py
@@ -31,7 +31,7 @@ def microsphere_affine_tube(F, f, kwargs, quadrature=BazantOh(n=21)):
 
 
 @displacement_pressure_split
-def microsphere_affine_force(x, fun, *args, **kwargs):
+def microsphere_affine_force(x, f, *args, **kwargs):
     """Micro-sphere model: Forces of affine stretch model as first Piola-
     Kirchhoff stress tensor embedded into a (u/p)-framework."""
 
@@ -55,7 +55,7 @@ def microsphere_affine_force(x, fun, *args, **kwargs):
     bulk = kwargs.pop("bulk")
 
     # fiber forces and state variable update
-    f, statevars = fun(lam, statevars_n, *args, **kwargs)
+    f, statevars = f(lam, statevars_n, *args, **kwargs)
 
     # Second Piola-Kirchhoff stress tensor
     SG = reshape(sum1(f / lam * sphere.weights * sphere.bases), 3, 3)
diff --git a/tests/test_microsphere.py b/tests/test_microsphere.py
index 11d7f72..9485113 100644
--- a/tests/test_microsphere.py
+++ b/tests/test_microsphere.py
@@ -19,7 +19,7 @@ def test_microsphere_force():
     umat = MaterialTensor(
         x=[F, affine.force.p, Zn],
         fun=affine.force,
-        kwargs={"fun": nh, "mu": 1.0, "bulk": 5000},
+        kwargs={"f": nh, "mu": 1.0, "bulk": 5000},
         statevars=1,
         triu=True,
     )
diff --git a/tests/test_template-materials.py b/tests/test_template-materials.py
index 0c4ff8e..f9f5539 100644
--- a/tests/test_template-materials.py
+++ b/tests/test_template-materials.py
@@ -2,7 +2,7 @@
 
 from matadi import MaterialTensorGeneral
 from matadi.models import NeoHookeOgdenRoxburgh, Morph, neo_hooke, volumetric
-from matadi.math import det, gradient
+from matadi.math import det, gradient, ones_like, zeros_like
 
 
 def fun(x, C10=0.5, bulk=5000):
@@ -12,13 +12,17 @@ def fun(x, C10=0.5, bulk=5000):
 
     W = neo_hooke(F, C10)
     U = volumetric(J, bulk)
+    
+    statevars_old = x[-1]
+    statevars_new = ones_like(statevars_old) # only for testing
+    statevars_new = zeros_like(statevars_old)
 
-    return gradient(W, F) + gradient(U, F), x[-1]
+    return gradient(W, F) + gradient(U, F), statevars_new
 
 
 def test_u_templates():
 
-    Custom = MaterialTensorGeneral(fun, nstatevars=1)
+    Custom = MaterialTensorGeneral(fun, statevars_shape=(1, 1))
 
     # Material as a function of `F` and `p`
     # with additional state variables `z`