diff --git a/src/vectise/colours.py b/src/vectise/colours.py
index 11690b3..b8bdd53 100644
--- a/src/vectise/colours.py
+++ b/src/vectise/colours.py
@@ -1,5 +1,7 @@
 from abc import ABC, abstractmethod
 import numpy as np
+from numpy.typing import ArrayLike
+
 # the following functions are taken from Ben Southgate:
 # https://bsouthga.dev/posts/colour-gradients-with-python
 
@@ -131,6 +133,7 @@ class LinearGradientColourMap(ColourMap):
         colours: list[str] | None = ["#ff0000", "#ffffff", "#0000ff"],
         min_value: float = 0,
         max_value: float = 1,
+        matrix: ArrayLike | None = None,
         bins: int = 100,
     ):
         self.colours = polylinear_gradient(colours, bins)
diff --git a/src/vectise/matrix.py b/src/vectise/matrix.py
index f395db3..35a1d1d 100644
--- a/src/vectise/matrix.py
+++ b/src/vectise/matrix.py
@@ -1,4 +1,5 @@
 import numpy as np
+from numpy.typing import ArrayLike
 import svg
 from .colours import ColourMap
 
@@ -6,8 +7,8 @@ from .colours import ColourMap
 class MatrixVisualisation:
     def __init__(
         self,
-        matrix: np.typing.NDArray,
         cmap: ColourMap,
+        matrix: ArrayLike,
         text: bool = False,
         labels: int | list[str] | bool = False,
     ):
diff --git a/src/vectise/neural_net.py b/src/vectise/neural_net.py
index 6cd3527..d9b35f6 100644
--- a/src/vectise/neural_net.py
+++ b/src/vectise/neural_net.py
@@ -3,12 +3,13 @@ import logging
 
 import itertools
 import numpy as np
+from numpy.typing import ArrayLike
 
 
 class NeuralNet:
     def __init__(
         self,
-        matrix: np.typing.NDArray | list[np.typing.NDArray] | None = None,
+        matrix: ArrayLike | list[ArrayLike] | None = None,
         shape: tuple[int] | None = None,
         node_colour: str = "blue",
         node_border: str = "black",
@@ -18,7 +19,7 @@ class NeuralNet:
 
         Parameters
         ----------
-        matrix : np.typing.NDArray | None
+        matrix : ArrayLike | None
             matrix or list of matrices
         shape : tuple[int] | None
             shape