diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc32106d..9c87e698 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,9 +10,9 @@ All notable changes to this project will be documented in this file. The format
 - Add optional keyword-arguments to `math.transpose(**kwargs)` to support optional `out` and `order`-keywords.
 - Add the attribute `RegionBoundary.tangents`, which contains a list of tangent unit vectors. For `quad` cell-types the length of this list is one and for `hexahedron` cell-types it is of length two.
 - Add `math.inplane(A, vectors)` to return the in-plane components of a symmetric tensor `A`, where the plane is defined by its standard unit vectors.
-- Add `constitution.autodiff.jax.Hyperelastic` as a feature-equivalent alternative to `Hyperelastic` with `jax` as backend.
-- Add `constitution.autodiff.jax.Material` as a feature-equivalent alternative to `MaterialAD` with `jax` as backend.
-- Add the MORPH-material formulation for a JAX-based material `felupe.constitution.autodiff.jax.models.lagrange.morph()`.
+- Add `constitution.jax.Hyperelastic` as a feature-equivalent alternative to `Hyperelastic` with `jax` as backend.
+- Add `constitution.jax.Material` as a feature-equivalent alternative to `MaterialAD` with `jax` as backend.
+- Add the MORPH-material formulation for a JAX-based material `felupe.constitution.jax.models.lagrange.morph()`.
 
 ### Changed
 - Change default `np.einsum(..., order="K")` to `np.einsum(..., order="C")` in the methods of `Field`, `FieldAxisymmetric`, `FieldPlaneStrain` and `FieldContainer`.
diff --git a/docs/felupe/constitution/jax.rst b/docs/felupe/constitution/jax.rst
index f79bc215..0024beae 100644
--- a/docs/felupe/constitution/jax.rst
+++ b/docs/felupe/constitution/jax.rst
@@ -11,33 +11,33 @@ This page contains material model formulations with automatic differentiation us
 
 .. autosummary::
    
-   constitution.autodiff.jax.Hyperelastic
-   constitution.autodiff.jax.Material
+   constitution.jax.Hyperelastic
+   constitution.jax.Material
 
-**Material Models for** :class:`felupe.constitution.autodiff.jax.Material`
+**Material Models for** :class:`felupe.constitution.jax.Material`
 
 .. autosummary::
 
-   felupe.constitution.autodiff.jax.models.lagrange.morph
+   felupe.constitution.jax.models.lagrange.morph
 
 **Tools**
 
 .. autosummary::
    
-   constitution.autodiff.jax.vmap
+   constitution.jax.vmap
 
 **Detailed API Reference**
 
-.. autoclass:: felupe.constitution.autodiff.jax.Hyperelastic
+.. autoclass:: felupe.constitution.jax.Hyperelastic
    :members:
    :undoc-members:
    :inherited-members:
 
-.. autoclass:: felupe.constitution.autodiff.jax.Material
+.. autoclass:: felupe.constitution.jax.Material
    :members:
    :undoc-members:
    :inherited-members:
 
-.. autofunction:: felupe.constitution.autodiff.jax.models.lagrange.morph
+.. autofunction:: felupe.constitution.jax.models.lagrange.morph
 
-.. autofunction:: felupe.constitution.autodiff.jax.vmap
+.. autofunction:: felupe.constitution.jax.vmap
diff --git a/docs/felupe/constitution/lagrange.rst b/docs/felupe/constitution/lagrange.rst
index 7a1de9e0..29470667 100644
--- a/docs/felupe/constitution/lagrange.rst
+++ b/docs/felupe/constitution/lagrange.rst
@@ -34,7 +34,7 @@ This page contains Total- and Updated-Lagrange material formulations with automa
    :undoc-members:
    :inherited-members:
 
-.. autoclass:: felupe.constitution.autodiff.tensortrax.Material
+.. autoclass:: felupe.constitution.tensortrax.Material
    :members:
    :undoc-members:
    :inherited-members:
diff --git a/src/felupe/constitution/__init__.py b/src/felupe/constitution/__init__.py
index fa97ca70..852c25a4 100644
--- a/src/felupe/constitution/__init__.py
+++ b/src/felupe/constitution/__init__.py
@@ -3,28 +3,6 @@
 from ._material import Material
 from ._mixed import NearlyIncompressible, ThreeFieldVariation
 from ._view import ViewMaterial, ViewMaterialIncompressible
-from .autodiff.tensortrax import Hyperelastic
-from .autodiff.tensortrax import Material as MaterialAD
-from .autodiff.tensortrax import total_lagrange, updated_lagrange
-from .autodiff.tensortrax.models.hyperelastic import (
-    alexander,
-    anssari_benam_bucchi,
-    arruda_boyce,
-    extended_tube,
-    finite_strain_viscoelastic,
-    isochoric_volumetric_split,
-    lopez_pamies,
-    miehe_goektepe_lulei,
-    mooney_rivlin,
-    neo_hooke,
-    ogden,
-    ogden_roxburgh,
-    saint_venant_kirchhoff,
-    third_order_deformation,
-    van_der_waals,
-    yeoh,
-)
-from .autodiff.tensortrax.models.lagrange import morph, morph_representative_directions
 from .hyperelasticity import NeoHooke, NeoHookeCompressible, OgdenRoxburgh, Volumetric
 from .linear_elasticity import (
     LinearElastic,
@@ -42,6 +20,28 @@
     linear_elastic,
     linear_elastic_plastic_isotropic_hardening,
 )
+from .tensortrax import Hyperelastic
+from .tensortrax import Material as MaterialAD
+from .tensortrax import total_lagrange, updated_lagrange
+from .tensortrax.models.hyperelastic import (
+    alexander,
+    anssari_benam_bucchi,
+    arruda_boyce,
+    extended_tube,
+    finite_strain_viscoelastic,
+    isochoric_volumetric_split,
+    lopez_pamies,
+    miehe_goektepe_lulei,
+    mooney_rivlin,
+    neo_hooke,
+    ogden,
+    ogden_roxburgh,
+    saint_venant_kirchhoff,
+    third_order_deformation,
+    van_der_waals,
+    yeoh,
+)
+from .tensortrax.models.lagrange import morph, morph_representative_directions
 
 __all__ = [
     "alexander",
diff --git a/src/felupe/constitution/autodiff/__init__.py b/src/felupe/constitution/autodiff/__init__.py
deleted file mode 100644
index 844e505e..00000000
--- a/src/felupe/constitution/autodiff/__init__.py
+++ /dev/null
@@ -1,10 +0,0 @@
-"""
-felupe.constitution.autodiff
-============================
-This module contains constitutive material classes and models for several backends. The
-primary used backend in FElupe is based on :mod:`tensortrax`.
-"""
-
-from . import jax, tensortrax
-
-__all__ = ["jax", "tensortrax"]
diff --git a/src/felupe/constitution/autodiff/jax/__init__.py b/src/felupe/constitution/jax/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/jax/__init__.py
rename to src/felupe/constitution/jax/__init__.py
diff --git a/src/felupe/constitution/autodiff/jax/_hyperelastic.py b/src/felupe/constitution/jax/_hyperelastic.py
similarity index 95%
rename from src/felupe/constitution/autodiff/jax/_hyperelastic.py
rename to src/felupe/constitution/jax/_hyperelastic.py
index b4a99527..2af44e63 100644
--- a/src/felupe/constitution/autodiff/jax/_hyperelastic.py
+++ b/src/felupe/constitution/jax/_hyperelastic.py
@@ -20,7 +20,7 @@
 
 import numpy as np
 
-from ..._material import Material
+from .._material import Material
 from ._tools import total_lagrange, vmap2
 
 
@@ -72,7 +72,7 @@ def neo_hooke(C, mu):
             "Strain energy function of the Neo-Hookean material formulation."
             return mu / 2 * (jnp.linalg.det(C) ** (-1/3) * jnp.trace(C) - 3)
 
-        umat = fem.constitution.autodiff.jax.Hyperelastic(neo_hooke, mu=1)
+        umat = fem.constitution.jax.Hyperelastic(neo_hooke, mu=1)
 
     and this code-block for material formulations with state variables.
 
@@ -83,6 +83,7 @@ def neo_hooke(C, mu):
     ..  code-block::
 
         import felupe as fem
+        import felupe.constitution.jax as mat
         import jax.numpy as np
 
         def viscoelastic(C, Cin, mu, eta, dtime):
@@ -101,9 +102,7 @@ def viscoelastic(C, Cin, mu, eta, dtime):
             # strain energy function and state variable
             return mu / 2 * (I1 - 3), Ci.ravel()
 
-        umat = fem.constitution.autodiff.jax.Hyperelastic(
-            viscoelastic, mu=1, eta=1, dtime=1, nstatevars=9
-        )
+        umat = mat.Hyperelastic(viscoelastic, mu=1, eta=1, dtime=1, nstatevars=9)
 
     ..  note::
         See the `documentation of JAX <https://jax.readthedocs.io>`_ for further
@@ -120,13 +119,14 @@ def viscoelastic(C, Cin, mu, eta, dtime):
         :context:
 
         >>> import felupe as fem
+        >>> import felupe.constitution.jax as mat
         >>> import jax.numpy as jnp
         >>>
         >>> def neo_hooke(C, mu):
         ...     "Strain energy function of the Neo-Hookean material formulation."
         ...     return mu / 2 * (jnp.linalg.det(C) ** (-1/3) * jnp.trace(C) - 3)
         >>>
-        >>> umat = fem.constitution.autodiff.jax.Hyperelastic(neo_hooke, mu=1)
+        >>> umat = mat.Hyperelastic(neo_hooke, mu=1)
         >>> ax = umat.plot(incompressible=True)
 
     ..  pyvista-plot::
diff --git a/src/felupe/constitution/autodiff/jax/_material.py b/src/felupe/constitution/jax/_material.py
similarity index 94%
rename from src/felupe/constitution/autodiff/jax/_material.py
rename to src/felupe/constitution/jax/_material.py
index 5fe604bf..6504d400 100644
--- a/src/felupe/constitution/autodiff/jax/_material.py
+++ b/src/felupe/constitution/jax/_material.py
@@ -20,7 +20,7 @@
 
 import numpy as np
 
-from ..._material import Material as MaterialDefault
+from .._material import Material as MaterialDefault
 from ._tools import vmap2
 
 
@@ -63,6 +63,7 @@ class Material(MaterialDefault):
     ..  code-block::
 
         import felupe as fem
+        import felupe.constitution.jax as mat
         import jax.numpy as jnp
 
         def neo_hooke(F, mu):
@@ -74,13 +75,14 @@ def neo_hooke(F, mu):
 
             return mu * F @ dev(Cu) @ jnp.linalg.inv(C)
 
-        umat = fem.constitution.autodiff.jax.Material(neo_hooke, mu=1)
+        umat = mat.Material(neo_hooke, mu=1)
 
     and this code-block for material formulations with state variables:
 
     ..  code-block::
 
         import felupe as fem
+        import felupe.constitution.jax as mat
         import jax.numpy as jnp
 
         def viscoelastic(F, Cin, mu, eta, dtime):
@@ -104,9 +106,7 @@ def viscoelastic(F, Cin, mu, eta, dtime):
             to_triu = lambda C: C[i, j]
             return F @ S, to_triu(Ci)
 
-        umat = fem.constitution.autodiff.jax.Material(
-            viscoelastic, mu=1, eta=1, dtime=1, nstatevars=6
-        )
+        umat = mat.Material(viscoelastic, mu=1, eta=1, dtime=1, nstatevars=6)
 
     ..  note::
         See the `documentation of JAX <https://jax.readthedocs.io>`_ for further
@@ -123,6 +123,7 @@ def viscoelastic(F, Cin, mu, eta, dtime):
         :context:
 
         >>> import felupe as fem
+        >>> import felupe.constitution.jax as mat
         >>> import jax.numpy as jnp
         >>>
         >>> def neo_hooke(F, mu):
@@ -134,7 +135,7 @@ def viscoelastic(F, Cin, mu, eta, dtime):
         ...
         ...     return mu * F @ dev(Cu) @ jnp.linalg.inv(C)
         >>>
-        >>> umat = fem.constitution.autodiff.jax.Material(neo_hooke, mu=1)
+        >>> umat = mat.Material(neo_hooke, mu=1)
         >>> ax = umat.plot(incompressible=True)
 
     ..  pyvista-plot::
diff --git a/src/felupe/constitution/autodiff/jax/_tools.py b/src/felupe/constitution/jax/_tools.py
similarity index 100%
rename from src/felupe/constitution/autodiff/jax/_tools.py
rename to src/felupe/constitution/jax/_tools.py
diff --git a/src/felupe/constitution/autodiff/jax/models/__init__.py b/src/felupe/constitution/jax/models/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/jax/models/__init__.py
rename to src/felupe/constitution/jax/models/__init__.py
diff --git a/src/felupe/constitution/autodiff/jax/models/hyperelastic/__init__.py b/src/felupe/constitution/jax/models/hyperelastic/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/jax/models/hyperelastic/__init__.py
rename to src/felupe/constitution/jax/models/hyperelastic/__init__.py
diff --git a/src/felupe/constitution/autodiff/jax/models/lagrange/__init__.py b/src/felupe/constitution/jax/models/lagrange/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/jax/models/lagrange/__init__.py
rename to src/felupe/constitution/jax/models/lagrange/__init__.py
diff --git a/src/felupe/constitution/autodiff/jax/models/lagrange/_morph.py b/src/felupe/constitution/jax/models/lagrange/_morph.py
similarity index 98%
rename from src/felupe/constitution/autodiff/jax/models/lagrange/_morph.py
rename to src/felupe/constitution/jax/models/lagrange/_morph.py
index 4166133f..b7e242f7 100644
--- a/src/felupe/constitution/autodiff/jax/models/lagrange/_morph.py
+++ b/src/felupe/constitution/jax/models/lagrange/_morph.py
@@ -122,9 +122,10 @@ def morph(F, statevars, p):
         :context:
 
         >>> import felupe as fem
+        >>> import felupe.constitution.jax as mat
         >>>
-        >>> umat = fem.constitution.autodiff.jax.Material(
-        ...     fem.constitution.autodiff.jax.models.lagrange.morph,
+        >>> umat = mat.Material(
+        ...     mat.models.lagrange.morph,
         ...     p=[0.039, 0.371, 0.174, 2.41, 0.0094, 6.84, 5.65, 0.244],
         ...     nstatevars=13,
         ... )
diff --git a/src/felupe/constitution/autodiff/tensortrax/__init__.py b/src/felupe/constitution/tensortrax/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/__init__.py
rename to src/felupe/constitution/tensortrax/__init__.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/_hyperelastic.py b/src/felupe/constitution/tensortrax/_hyperelastic.py
similarity index 98%
rename from src/felupe/constitution/autodiff/tensortrax/_hyperelastic.py
rename to src/felupe/constitution/tensortrax/_hyperelastic.py
index 6687034d..95ed7a85 100644
--- a/src/felupe/constitution/autodiff/tensortrax/_hyperelastic.py
+++ b/src/felupe/constitution/tensortrax/_hyperelastic.py
@@ -19,8 +19,8 @@
 import numpy as np
 import tensortrax as tr
 
-from ....math import cdya_ik, dot, transpose
-from ..._material import Material
+from ...math import cdya_ik, dot, transpose
+from .._material import Material
 
 
 class Hyperelastic(Material):
diff --git a/src/felupe/constitution/autodiff/tensortrax/_material.py b/src/felupe/constitution/tensortrax/_material.py
similarity index 93%
rename from src/felupe/constitution/autodiff/tensortrax/_material.py
rename to src/felupe/constitution/tensortrax/_material.py
index 3592404f..4b6e2368 100644
--- a/src/felupe/constitution/autodiff/tensortrax/_material.py
+++ b/src/felupe/constitution/tensortrax/_material.py
@@ -19,7 +19,7 @@
 import numpy as np
 import tensortrax as tr
 
-from ..._material import Material as MaterialDefault
+from .._material import Material as MaterialDefault
 
 
 class Material(MaterialDefault):
@@ -66,6 +66,7 @@ class Material(MaterialDefault):
     ..  code-block::
 
         import felupe as fem
+        import felupe.constitution.tensortrax as mat
         import tensortrax.math as tm
 
         def neo_hooke(F, mu):
@@ -76,13 +77,14 @@ def neo_hooke(F, mu):
 
             return mu * F @ tm.special.dev(Cu) @ tm.linalg.inv(C)
 
-        umat = fem.MaterialAD(neo_hooke, mu=1)
+        umat = mat.Material(neo_hooke, mu=1)
 
     and this code-block for material formulations with state variables:
 
     ..  code-block::
 
         import felupe as fem
+        import felupe.constitution.tensortrax as mat
         import tensortrax.math as tm
 
         def viscoelastic(F, Cin, mu, eta, dtime):
@@ -102,9 +104,7 @@ def viscoelastic(F, Cin, mu, eta, dtime):
             # first Piola-Kirchhoff stress tensor and state variable
             return F @ S, tm.special.triu_1d(Ci)
 
-        umat = fem.MaterialAD(
-            viscoelastic, mu=1, eta=1, dtime=1, nstatevars=6
-        )
+        umat = mat.Material(viscoelastic, mu=1, eta=1, dtime=1, nstatevars=6)
 
     ..  note::
         See the `documentation of tensortrax <https://github.com/adtzlr/tensortrax>`_
@@ -116,6 +116,7 @@ def viscoelastic(F, Cin, mu, eta, dtime):
         :context:
 
         >>> import felupe as fem
+        >>> import felupe.constitution.tensortrax as mat
         >>> import tensortrax.math as tm
         >>>
         >>> def neo_hooke(F, mu):
@@ -123,7 +124,7 @@ def viscoelastic(F, Cin, mu, eta, dtime):
         ...     S = mu * tm.special.dev(tm.linalg.det(C)**(-1/3) * C) @ tm.linalg.inv(C)
         ...     return F @ S
         >>>
-        >>> umat = fem.MaterialAD(neo_hooke, mu=1)
+        >>> umat = mat.Material(neo_hooke, mu=1)
         >>> ax = umat.plot(incompressible=True)
 
     ..  pyvista-plot::
diff --git a/src/felupe/constitution/autodiff/tensortrax/_total_lagrange.py b/src/felupe/constitution/tensortrax/_total_lagrange.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/_total_lagrange.py
rename to src/felupe/constitution/tensortrax/_total_lagrange.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/_updated_lagrange.py b/src/felupe/constitution/tensortrax/_updated_lagrange.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/_updated_lagrange.py
rename to src/felupe/constitution/tensortrax/_updated_lagrange.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/__init__.py b/src/felupe/constitution/tensortrax/models/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/__init__.py
rename to src/felupe/constitution/tensortrax/models/__init__.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/__init__.py b/src/felupe/constitution/tensortrax/models/hyperelastic/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/__init__.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/__init__.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_alexander.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_alexander.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_alexander.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_alexander.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_anssari_benam_bucchi.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_anssari_benam_bucchi.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_anssari_benam_bucchi.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_anssari_benam_bucchi.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_arruda_boyce.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_arruda_boyce.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_arruda_boyce.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_arruda_boyce.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_extended_tube.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_extended_tube.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_extended_tube.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_extended_tube.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_finite_strain_viscoelastic.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_finite_strain_viscoelastic.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_finite_strain_viscoelastic.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_finite_strain_viscoelastic.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_helpers.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_helpers.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_helpers.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_helpers.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_lopez_pamies.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_lopez_pamies.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_lopez_pamies.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_lopez_pamies.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_miehe_goektepe_lulei.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_miehe_goektepe_lulei.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_miehe_goektepe_lulei.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_miehe_goektepe_lulei.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_mooney_rivlin.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_mooney_rivlin.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_mooney_rivlin.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_mooney_rivlin.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_morph_representative_directions.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_morph_representative_directions.py
similarity index 98%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_morph_representative_directions.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_morph_representative_directions.py
index 064dfc8b..c9ca837e 100644
--- a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_morph_representative_directions.py
+++ b/src/felupe/constitution/tensortrax/models/hyperelastic/_morph_representative_directions.py
@@ -43,7 +43,7 @@ def morph_representative_directions(C, statevars, p, ε=1e-8):
         :context:
 
         >>> import felupe as fem
-        >>> import felupe.constitution.autodiff.tensortrax.models as models
+        >>> import felupe.constitution.tensortrax.models as models
         >>>
         >>> umat = fem.Hyperelastic(
         ...     models.hyperelastic.morph_representative_directions,
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_neo_hooke.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_neo_hooke.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_neo_hooke.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_neo_hooke.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_ogden.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_ogden.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_ogden.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_ogden.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_ogden_roxburgh.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_ogden_roxburgh.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_ogden_roxburgh.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_ogden_roxburgh.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_saint_venant_kirchhoff.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_saint_venant_kirchhoff.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_saint_venant_kirchhoff.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_saint_venant_kirchhoff.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_third_order_deformation.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_third_order_deformation.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_third_order_deformation.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_third_order_deformation.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_van_der_waals.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_van_der_waals.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_van_der_waals.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_van_der_waals.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_yeoh.py b/src/felupe/constitution/tensortrax/models/hyperelastic/_yeoh.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/_yeoh.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/_yeoh.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/__init__.py b/src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/__init__.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/__init__.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_chain.py b/src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_chain.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_chain.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_chain.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_framework_affine.py b/src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_framework_affine.py
similarity index 97%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_framework_affine.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_framework_affine.py
index ae3053f4..2afe444e 100644
--- a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_framework_affine.py
+++ b/src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_framework_affine.py
@@ -1,7 +1,7 @@
 from tensortrax.math import einsum, sqrt
 from tensortrax.math.linalg import det, inv
 
-from .......quadrature import BazantOh
+from ......quadrature import BazantOh
 
 
 def affine_stretch(C, f, kwargs, quadrature=BazantOh(n=21)):
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_framework_nonaffine.py b/src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_framework_nonaffine.py
similarity index 95%
rename from src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_framework_nonaffine.py
rename to src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_framework_nonaffine.py
index 7fd122e6..7f781324 100644
--- a/src/felupe/constitution/autodiff/tensortrax/models/hyperelastic/microsphere/_framework_nonaffine.py
+++ b/src/felupe/constitution/tensortrax/models/hyperelastic/microsphere/_framework_nonaffine.py
@@ -1,7 +1,7 @@
 from tensortrax.math import einsum, sqrt
 from tensortrax.math.linalg import det, inv
 
-from .......quadrature import BazantOh
+from ......quadrature import BazantOh
 
 
 def nonaffine_stretch(C, p, f, kwargs, quadrature=BazantOh(n=21)):
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/__init__.py b/src/felupe/constitution/tensortrax/models/lagrange/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/lagrange/__init__.py
rename to src/felupe/constitution/tensortrax/models/lagrange/__init__.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/_morph.py b/src/felupe/constitution/tensortrax/models/lagrange/_morph.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/lagrange/_morph.py
rename to src/felupe/constitution/tensortrax/models/lagrange/_morph.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/_morph_representative_directions.py b/src/felupe/constitution/tensortrax/models/lagrange/_morph_representative_directions.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/lagrange/_morph_representative_directions.py
rename to src/felupe/constitution/tensortrax/models/lagrange/_morph_representative_directions.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/_morph_uniaxial.py b/src/felupe/constitution/tensortrax/models/lagrange/_morph_uniaxial.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/lagrange/_morph_uniaxial.py
rename to src/felupe/constitution/tensortrax/models/lagrange/_morph_uniaxial.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/microsphere/__init__.py b/src/felupe/constitution/tensortrax/models/lagrange/microsphere/__init__.py
similarity index 100%
rename from src/felupe/constitution/autodiff/tensortrax/models/lagrange/microsphere/__init__.py
rename to src/felupe/constitution/tensortrax/models/lagrange/microsphere/__init__.py
diff --git a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/microsphere/_framework_affine.py b/src/felupe/constitution/tensortrax/models/lagrange/microsphere/_framework_affine.py
similarity index 95%
rename from src/felupe/constitution/autodiff/tensortrax/models/lagrange/microsphere/_framework_affine.py
rename to src/felupe/constitution/tensortrax/models/lagrange/microsphere/_framework_affine.py
index b9ae90bb..1864e5e5 100644
--- a/src/felupe/constitution/autodiff/tensortrax/models/lagrange/microsphere/_framework_affine.py
+++ b/src/felupe/constitution/tensortrax/models/lagrange/microsphere/_framework_affine.py
@@ -1,7 +1,7 @@
 from tensortrax.math import einsum, sqrt, trace
 from tensortrax.math.linalg import det, inv
 
-from .......quadrature import BazantOh
+from ......quadrature import BazantOh
 from ...._total_lagrange import total_lagrange
 
 
diff --git a/tests/test_constitution.py b/tests/test_constitution.py
index 94f7a5e8..4e436101 100644
--- a/tests/test_constitution.py
+++ b/tests/test_constitution.py
@@ -29,7 +29,7 @@
 import pytest
 
 import felupe as fem
-import felupe.constitution.autodiff.tensortrax.models as models
+import felupe.constitution.tensortrax.models as models
 
 
 def pre(sym, add_identity, add_random=False):
diff --git a/tests/test_constitution_jax.py b/tests/test_constitution_jax.py
index d58bd901..97321173 100644
--- a/tests/test_constitution_jax.py
+++ b/tests/test_constitution_jax.py
@@ -28,6 +28,7 @@
 import pytest
 
 import felupe as fem
+import felupe.constitution.jax as mat
 
 
 def test_vmap():
@@ -39,8 +40,8 @@ def f(x, a=1.0):
         def g(x, y, a=1.0, **kwargs):
             return x
 
-        vf = fem.constitution.autodiff.jax.vmap(f)
-        vg = fem.constitution.autodiff.jax.vmap(g)
+        vf = fem.constitution.jax.vmap(f)
+        vg = fem.constitution.jax.vmap(g)
 
         x = np.eye(3).reshape(1, 3, 3) * np.ones((10, 1, 1))
 
@@ -69,10 +70,8 @@ def W(C, C10, K):
             I1 = I3 ** (-1 / 3) * jnp.trace(C)
             return C10 * (I1 - 3) + K * (J - 1) ** 2 / 2
 
-        umat = fem.constitution.autodiff.jax.Hyperelastic(
-            W, C10=0.5, K=2.0, parallel=True
-        )
-        umat = fem.constitution.autodiff.jax.Hyperelastic(W, C10=0.5, K=2.0, jit=True)
+        umat = fem.constitution.jax.Hyperelastic(W, C10=0.5, K=2.0, parallel=True)
+        umat = mat.Hyperelastic(W, C10=0.5, K=2.0, jit=True)
         mesh = fem.Cube(n=2)
         region = fem.RegionHexahedron(mesh)
         field = fem.FieldContainer([fem.Field(region, dim=3)])
@@ -103,9 +102,7 @@ def W(C, statevars, C10, K):
 
         W.kwargs = {"C10": 0.5}
 
-        umat = fem.constitution.autodiff.jax.Hyperelastic(
-            W, C10=0.5, K=2.0, nstatevars=1, jit=True
-        )
+        umat = mat.Hyperelastic(W, C10=0.5, K=2.0, nstatevars=1, jit=True)
         mesh = fem.Cube(n=2)
         region = fem.RegionHexahedron(mesh)
         field = fem.FieldContainer([fem.Field(region, dim=3)])
@@ -136,10 +133,8 @@ def dWdF(F, C10, K):
             P = 2 * C10 * F @ dev(Cu) @ jnp.linalg.inv(C)
             return P + K * (J - 1) * J * jnp.linalg.inv(C)
 
-        umat = fem.constitution.autodiff.jax.Material(
-            dWdF, C10=0.5, K=2.0, parallel=True
-        )
-        umat = fem.constitution.autodiff.jax.Material(dWdF, C10=0.5, K=2.0, jit=True)
+        umat = mat.Material(dWdF, C10=0.5, K=2.0, parallel=True)
+        umat = fem.constitution.jax.Material(dWdF, C10=0.5, K=2.0, jit=True)
         mesh = fem.Cube(n=2)
         region = fem.RegionHexahedron(mesh)
         field = fem.FieldContainer([fem.Field(region, dim=3)])
@@ -173,9 +168,7 @@ def dWdF(F, statevars, C10, K):
 
         dWdF.kwargs = {"C10": 0.5}
 
-        umat = fem.constitution.autodiff.jax.Material(
-            dWdF, C10=0.5, K=2.0, nstatevars=1, jit=True
-        )
+        umat = mat.Material(dWdF, C10=0.5, K=2.0, nstatevars=1, jit=True)
         mesh = fem.Cube(n=2)
         region = fem.RegionHexahedron(mesh)
         field = fem.FieldContainer([fem.Field(region, dim=3)])
@@ -195,8 +188,8 @@ def dWdF(F, statevars, C10, K):
 
 def test_material_included_jax_statevars():
     try:
-        umat = fem.constitution.autodiff.jax.Material(
-            fem.constitution.autodiff.jax.models.lagrange.morph,
+        umat = mat.Material(
+            fem.constitution.jax.models.lagrange.morph,
             p=[0.039, 0.371, 0.174, 2.41, 0.0094, 6.84, 5.65, 0.244],
             nstatevars=13,
         )