diff --git a/src/main/java/qupath/bioimageio/spec/tensor/axes/BatchAxis.java b/src/main/java/qupath/bioimageio/spec/tensor/axes/BatchAxis.java
index d448b00..2b0ec25 100644
--- a/src/main/java/qupath/bioimageio/spec/tensor/axes/BatchAxis.java
+++ b/src/main/java/qupath/bioimageio/spec/tensor/axes/BatchAxis.java
@@ -23,9 +23,9 @@
 import java.util.List;
 
 /**
- * A base Batch class.
+ * An axis that relates to data batches.
  */
-class BatchAxis extends AxisBase {
+public class BatchAxis extends AxisBase {
     private final int size;
     private final boolean concatenable = true;
 
@@ -49,4 +49,3 @@ public void validate(List<? extends BaseTensor> tensors) {
         // fixed size doesn't need validation
     }
 }
-
diff --git a/src/main/java/qupath/bioimageio/spec/tensor/axes/ChannelAxis.java b/src/main/java/qupath/bioimageio/spec/tensor/axes/ChannelAxis.java
index a399704..fd78100 100644
--- a/src/main/java/qupath/bioimageio/spec/tensor/axes/ChannelAxis.java
+++ b/src/main/java/qupath/bioimageio/spec/tensor/axes/ChannelAxis.java
@@ -23,9 +23,9 @@
 import java.util.List;
 
 /**
- * A list of channels.
+ * An axis corresponding to a set of channels.
  */
-class ChannelAxis extends AxisBase implements ScaledAxis {
+public class ChannelAxis extends AxisBase implements ScaledAxis {
     private final List<String> channel_names;
 
     ChannelAxis(String id, String description, List<String> channel_names) {
@@ -53,4 +53,3 @@ public double getScale() {
         return 1;
     }
 }
-
diff --git a/src/main/java/qupath/bioimageio/spec/tensor/axes/IndexAxes.java b/src/main/java/qupath/bioimageio/spec/tensor/axes/IndexAxes.java
index 6ded19a..b73073d 100644
--- a/src/main/java/qupath/bioimageio/spec/tensor/axes/IndexAxes.java
+++ b/src/main/java/qupath/bioimageio/spec/tensor/axes/IndexAxes.java
@@ -25,12 +25,18 @@
  * Axes that relate to indexes rather than space/time/channel/batch, for example lists of objects.
  */
 public class IndexAxes {
-    abstract static class IndexAxisBase extends AxisBase implements ScaledAxis {
+    /**
+     * An axis relating to indexes, for example a list of objects.
+     */
+    public static class IndexAxis extends AxisBase implements ScaledAxis {
+        private final Size size;
         private final double scale = 1.0;
         private final String unit = null;
+        private final boolean concatenable = false;
 
-        IndexAxisBase(String id, String description) {
+        IndexAxis(String id, String description, Size size) {
             super(id, description);
+            this.size = size;
         }
 
         @Override
@@ -42,16 +48,7 @@ public AxisType getType() {
         public double getScale() {
             return 1;
         }
-    }
 
-    static class IndexAxis extends IndexAxisBase {
-        private final Size size;
-        private final boolean concatenable = false;
-
-        IndexAxis(String id, String description, Size size) {
-            super(id, description);
-            this.size = size;
-        }
         @Override
         public Size getSize() {
             return this.size;
@@ -61,6 +58,7 @@ public Size getSize() {
         public void validate(List<? extends BaseTensor> tensors) {
             size.validate(tensors);
         }
+
     }
 
 }
diff --git a/src/main/java/qupath/bioimageio/spec/tensor/axes/SpaceAxes.java b/src/main/java/qupath/bioimageio/spec/tensor/axes/SpaceAxes.java
index 7fecb46..bc95047 100644
--- a/src/main/java/qupath/bioimageio/spec/tensor/axes/SpaceAxes.java
+++ b/src/main/java/qupath/bioimageio/spec/tensor/axes/SpaceAxes.java
@@ -80,18 +80,24 @@ public static SpaceUnit getUnit(String unit) {
         }
     }
 
-    abstract static class SpaceAxisBase extends AxisBase implements ScaledAxis {
+    /**
+     * An axis that relates to physical space (X, Y, Z)
+     */
+    public static class SpaceAxis extends AxisBase implements ScaledAxis {
         private final SpaceUnit unit;
         private final double scale;
+        private final Size size;
+        private final boolean concatenable = false;
 
-        SpaceAxisBase(String id, String description, String unit, double scale) {
-            this(id, description, SpaceUnit.valueOf(unit.toUpperCase()), scale);
+        SpaceAxis(String id, String description, String unit, double scale, Size size) {
+            this(id, description, SpaceUnit.valueOf(unit.toUpperCase()), scale, size);
         }
 
-        SpaceAxisBase(String id, String description, SpaceUnit unit, double scale) {
+        SpaceAxis(String id, String description, SpaceUnit unit, double scale, Size size) {
             super(id, description);
             this.unit = unit;
             this.scale = scale;
+            this.size = size;
         }
 
         @Override
@@ -104,33 +110,25 @@ public double getScale() {
             return this.scale;
         }
 
+        /**
+         * Gets the unit for the space dimension
+         * @return the unit (hopefully SI, but maybe imperial)
+         */
         public SpaceUnit getUnit() {
             return unit;
         }
-    }
-
-    static class SpaceAxis extends SpaceAxisBase {
-        private final Size size;
-        private final boolean concatenable = false;
 
-        SpaceAxis(String id, String description, String unit, double scale, Size size) {
-            super(id, description, unit.isEmpty() ? "NO_UNIT" : unit, scale);
-            this.size = size;
+        @Override
+        public void validate(List<? extends BaseTensor> tensors) {
+            getSize().validate(tensors);
         }
 
         @Override
         public Size getSize() {
             return this.size;
         }
-
-        @Override
-        public void validate(List<? extends BaseTensor> tensors) {
-            getSize().validate(tensors);
-        }
     }
 
-
-
     static class SpaceAxisWithHalo extends SpaceAxis implements WithHalo {
         private ReferencedSize size;
         private final int halo;
diff --git a/src/main/java/qupath/bioimageio/spec/tensor/axes/TimeAxes.java b/src/main/java/qupath/bioimageio/spec/tensor/axes/TimeAxes.java
index a5ffdbe..b95b409 100644
--- a/src/main/java/qupath/bioimageio/spec/tensor/axes/TimeAxes.java
+++ b/src/main/java/qupath/bioimageio/spec/tensor/axes/TimeAxes.java
@@ -26,16 +26,23 @@
  */
 public class TimeAxes {
 
-    abstract static class TimeAxisBase extends AxisBase implements ScaledAxis {
+    /**
+     * An axis that relates to physical time.
+     */
+    public static class TimeAxis extends AxisBase implements ScaledAxis {
         private final String type = "time";
         private final TimeUnit unit;
         private final double scale;
 
-        TimeAxisBase(String id, String description, TimeUnit unit, double scale) {
+        private final Size size;
+
+        TimeAxis(String id, String description, TimeUnit unit, double scale, Size size) {
             super(id, description);
             this.unit = unit;
             this.scale = scale;
+            this.size = size;
         }
+
         @Override
         public AxisType getType() {
             return AxisType.T;
@@ -46,19 +53,6 @@ public double getScale() {
             return scale;
         }
 
-        public TimeUnit getUnit() {
-            return unit;
-        }
-    }
-
-    static class TimeAxis extends TimeAxisBase {
-        private final Size size;
-
-        TimeAxis(String id, String description, TimeUnit unit, double scale, Size size) {
-            super(id, description, unit, scale);
-            this.size = size;
-        }
-
         @Override
         public Size getSize() {
             return size;
@@ -68,6 +62,15 @@ public Size getSize() {
         public void validate(List<? extends BaseTensor> tensors) {
             getSize().validate(tensors);
         }
+
+        /**
+         * Gets the unit for this axis.
+         * @return the unit of time, hopefully SI but maybe something strange like "Day"
+         */
+        public TimeUnit getUnit() {
+            return unit;
+        }
+
     }
 
     static class TimeAxisWithHalo extends TimeAxis implements WithHalo {