diff --git a/arrayfire/ArrayFire.java b/arrayfire/ArrayFire.java
index bbb9474..bae2b31 100644
--- a/arrayfire/ArrayFire.java
+++ b/arrayfire/ArrayFire.java
@@ -99,24 +99,24 @@ public static Scope scope() {
     /**
      * Sorts a tensor over D0.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<DT, D0, D1, D2, D3> sort(
-        Tensor<DT, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sort(Tensor<T, S> tensor) {
         return sort(tensor, D0);
     }
 
     /**
      * Sorts a tensor over the given dimension.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<DT, D0, D1, D2, D3> sort(
-        Tensor<DT, D0, D1, D2, D3> tensor, Dim dim) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sort(Tensor<T, S> tensor,
+                                                                                            Dim dim) {
         return sort(tensor, dim, true);
     }
 
     /**
      * Sorts a tensor over the given dimension in ascending or descending order.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<DT, D0, D1, D2, D3> sort(
-        Tensor<DT, D0, D1, D2, D3> tensor, Dim dim, boolean ascending) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sort(Tensor<T, S> tensor,
+                                                                                            Dim dim,
+                                                                                            boolean ascending) {
         return operation("sort")
                    .inputs(tensor)
                    .outputs(tensor.prototype())
@@ -127,40 +127,39 @@ public static Scope scope() {
     /**
      * Returns a prototype tensor with the given type and shape.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Prototype<T, D0, D1, D2, D3> prototype(
-        T type, Shape<D0, D1, D2, D3> shape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Prototype<T, S> prototype(T type, S shape) {
         return new Prototype<>(type, shape);
     }
 
     /**
      * Returns a prototype tensor with the same type and shape as the given tensor.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Prototype<T, D0, D1, D2, D3> prototype(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Prototype<T, S> prototype(
+        Tensor<T, S> tensor) {
         return new Prototype<>(tensor.type(), tensor.shape());
     }
 
     /**
      * Sorts a tensor over D0 and returns the values and indices of original values.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> SortIndexResult<DT, D0, D1, D2, D3> sortIndex(
-        Tensor<DT, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> SortIndexResult<T, S> sortIndex(
+        Tensor<T, S> tensor) {
         return sortIndex(tensor, D0);
     }
 
     /**
      * Sorts a tensor over the given dimension and returns the values and indices of original values.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> SortIndexResult<DT, D0, D1, D2, D3> sortIndex(
-        Tensor<DT, D0, D1, D2, D3> tensor, Dim dim) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> SortIndexResult<T, S> sortIndex(
+        Tensor<T, S> tensor, Dim dim) {
         return sortIndex(tensor, dim, true);
     }
 
     /**
      * Sorts a tensor over the given dimension in ascending or descending order and returns the values and indices of original values.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> SortIndexResult<DT, D0, D1, D2, D3> sortIndex(
-        Tensor<DT, D0, D1, D2, D3> tensor, Dim dim, boolean ascending) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> SortIndexResult<T, S> sortIndex(
+        Tensor<T, S> tensor, Dim dim, boolean ascending) {
         var pair = operation("sort_index")
                        .inputs(tensor)
                        .outputs(prototype(tensor.type(), tensor.shape()), prototype(U32, tensor.shape()))
@@ -182,7 +181,7 @@ public static <D extends Num<D>> Index<D> permutation(D dim) {
     /**
      * Creates a device tensor from the given native array.
      */
-    public static <DT extends DataType<?, ?>, AT extends NativeArray<DT, ?, ?>> Tensor<DT, N, U, U, U> create(
+    public static <DT extends DataType<?, ?>, AT extends NativeArray<DT, ?, ?>> Tensor<DT, Shape<N, U, U, U>> create(
         AT array) {
         return create(array, shape(n(array.length())));
     }
@@ -190,8 +189,8 @@ public static <D extends Num<D>> Index<D> permutation(D dim) {
     /**
      * Creates a device tensor from the given native array and shape.
      */
-    public static <DT extends DataType<?, ?>, AT extends NativeArray<DT, ?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<DT, D0, D1, D2, D3> create(
-        AT array, Shape<D0, D1, D2, D3> shape) {
+    public static <DT extends DataType<?, ?>, AT extends NativeArray<DT, ?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<DT, S> create(
+        AT array, S shape) {
         return operation("create")
                    .inputs()
                    .outputs(prototype(array.type(), shape))
@@ -206,7 +205,7 @@ public static <D extends Num<D>> Index<D> permutation(D dim) {
      * This is not recommended in a production setting, as memory will be copied twice. Instead, use {@link #create(NativeArray)}.
      */
     @SafeVarargs
-    public static <JT, AT extends NativeArray<DT, JT, ?>, DT extends DataType<AT, ?>> Tensor<DT, N, U, U, U> create(
+    public static <JT, AT extends NativeArray<DT, JT, ?>, DT extends DataType<AT, ?>> Tensor<DT, Shape<N, U, U, U>> create(
         DT type, JT... values) {
         return tidy(() -> {
             var array = type.create(values.length);
@@ -222,7 +221,7 @@ public static <D extends Num<D>> Index<D> permutation(D dim) {
      * This is not recommended in a production setting, as memory will be copied twice. Instead, use {@link #create(NativeArray)}.
      */
     @SuppressWarnings("unchecked")
-    public static <JT, JTA, AT extends NativeArray<DT, JT, JTA>, DT extends DataType<AT, ?>> Tensor<DT, N, U, U, U> create(
+    public static <JT, JTA, AT extends NativeArray<DT, JT, JTA>, DT extends DataType<AT, ?>> Tensor<DT, Shape<N, U, U, U>> create(
         DT type, JTA values) {
         return tidy(() -> {
             var length = Array.getLength(values);
@@ -237,50 +236,50 @@ public static <D extends Num<D>> Index<D> permutation(D dim) {
     /**
      * Creates a {@link F32} device tensor from the given float values.
      */
-    public static Tensor<F32, N, U, U, U> create(float... values) {
+    public static Tensor<F32, Shape<N, U, U, U>> create(float... values) {
         return create(F32, values);
     }
 
     /**
      * Creates a {@link F64} device tensor from the given double values.
      */
-    public static Tensor<F64, N, U, U, U> create(double... values) {
+    public static Tensor<F64, Shape<N, U, U, U>> create(double... values) {
         return create(F64, values);
     }
 
     /**
      * Creates a {@link S32} device tensor from the given byte values.
      */
-    public static Tensor<S32, N, U, U, U> create(int... values) {
+    public static Tensor<S32, Shape<N, U, U, U>> create(int... values) {
         return create(S32, values);
     }
 
     /**
      * Creates a constant scalar {@link F32} device tensor from the given float value.
      */
-    public static Tensor<F32, U, U, U, U> constant(float value) {
+    public static Tensor<F32, Shape<U, U, U, U>> constant(float value) {
         return constant(F32, value);
     }
 
     /**
      * Creates a constant scalar {@link F64} device tensor from the given float value.
      */
-    public static Tensor<F64, U, U, U, U> constant(double value) {
+    public static Tensor<F64, Shape<U, U, U, U>> constant(double value) {
         return constant(F64, value);
     }
 
     /**
      * Creates a constant scalar device tensor from the given type and double value.
      */
-    public static <DT extends DataType<?, ?>> Tensor<DT, U, U, U, U> constant(DT type, double value) {
+    public static <DT extends DataType<?, ?>> Tensor<DT, Shape<U, U, U, U>> constant(DT type, double value) {
         return constant(type, shape(u()), value);
     }
 
     /**
      * Creates a constant device tensor from the given type, shape, and double value.
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<DT, D0, D1, D2, D3> constant(
-        DT type, Shape<D0, D1, D2, D3> shape, double value) {
+    public static <DT extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<DT, S> constant(DT type, S shape,
+                                                                                                  double value) {
         return operation("constant")
                    .inputs()
                    .outputs(prototype(type, shape))
@@ -314,8 +313,8 @@ public static Index<N> seq(int begin, int endInclusive) {
     /**
      * Returns a lookup index using the given tensor as lookup values (indices).
      */
-    public static <DT extends DataType<?, ?>, D0 extends Num<D0>> Index<D0> seq(Tensor<DT, D0, U, U, U> index) {
-        return new Index<>(index, index.d0()::create);
+    public static <DT extends DataType<?, ?>, D0 extends Num<D0>> Index<D0> seq(Tensor<DT, Shape<D0, U, U, U>> index) {
+        return new Index<>(index, index.shape().d0()::create);
     }
 
     /**
@@ -349,15 +348,15 @@ public static Shape<N, U, U, U> shape(int d0) {
     /**
      * Returns a 1D shape of the given dimension.
      */
-    public static <D0 extends Num<?>> Shape<D0, U, U, U> shape(D0 d0) {
+    public static <D0 extends Num<D0>> Shape<D0, U, U, U> shape(D0 d0) {
         return new Shape<>(d0, u(), u(), u());
     }
 
-    public static <D0 extends Num<?>> Shape<D0, N, U, U> shape(D0 d0, int d1) {
+    public static <D0 extends Num<D0>> Shape<D0, N, U, U> shape(D0 d0, int d1) {
         return new Shape<>(d0, n(d1), u(), u());
     }
 
-    public static <D1 extends Num<?>> Shape<N, D1, U, U> shape(int d0, D1 d1) {
+    public static <D1 extends Num<D1>> Shape<N, D1, U, U> shape(int d0, D1 d1) {
         return new Shape<>(n(d0), d1, u(), u());
     }
 
@@ -365,7 +364,7 @@ public static Shape<N, N, U, U> shape(int d0, int d1) {
         return new Shape<>(n(d0), n(d1), u(), u());
     }
 
-    public static <D0 extends Num<?>, D1 extends Num<?>> Shape<D0, D1, U, U> shape(D0 d0, D1 d1) {
+    public static <D0 extends Num<D0>, D1 extends Num<D1>> Shape<D0, D1, U, U> shape(D0 d0, D1 d1) {
         return new Shape<>(d0, d1, u(), u());
     }
 
@@ -377,50 +376,50 @@ public static Shape<N, N, N, N> shape(int d0, int d1, int d2, int d3) {
         return new Shape<>(n(d0), n(d1), n(d2), n(d3));
     }
 
-    public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>> Shape<D0, D1, D2, U> shape(D0 d0, D1 d1,
-                                                                                                       D2 d2) {
+    public static <D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>> Shape<D0, D1, D2, U> shape(D0 d0, D1 d1,
+                                                                                                          D2 d2) {
         return new Shape<>(d0, d1, d2, u());
     }
 
-    public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Shape<D0, D1, D2, D3> shape(
+    public static <D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Shape<D0, D1, D2, D3> shape(
         D0 d0, D1 d1, D2 d2, D3 d3) {
         return new Shape<>(d0, d1, d2, d3);
     }
 
-    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Operation.Builder.Unary<IT, D0, D1, D2, D3>.Single<T, U, D1, D2, D3> reduce(
-        String name, Tensor<IT, D0, D1, D2, D3> a,
+    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Operation.Builder.Unary<Tensor<IT, Shape<D0, D1, D2, D3>>>.Single<Tensor<T, Shape<U, D1, D2, D3>>> reduce(
+        String name, Tensor<IT, Shape<D0, D1, D2, D3>> a,
         Functions.Function3<MemorySegment, MemorySegment, Integer, Integer> method, arrayfire.D0 dim, T resultType) {
         return operation(name)
                    .inputs(a)
-                   .outputs(prototype(resultType, shape(u(), a.d1(), a.d2(), a.d3())))
+                   .outputs(prototype(resultType, shape(u(), a.shape().d1(), a.shape().d2(), a.shape().d3())))
                    .operation(ptr -> method.apply(ptr, a.dereference(), dim.index()));
     }
 
-    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Operation.Builder.Unary<IT, D0, D1, D2, D3>.Single<T, D0, U, D2, D3> reduce(
-        String name, Tensor<IT, D0, D1, D2, D3> a,
+    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Operation.Builder.Unary<Tensor<IT, Shape<D0, D1, D2, D3>>>.Single<Tensor<T, Shape<D0, U, D2, D3>>> reduce(
+        String name, Tensor<IT, Shape<D0, D1, D2, D3>> a,
         Functions.Function3<MemorySegment, MemorySegment, Integer, Integer> method, arrayfire.D1 dim, T resultType) {
         return operation(name)
                    .inputs(a)
-                   .outputs(prototype(resultType, shape(a.d0(), u(), a.d2(), a.d3())))
+                   .outputs(prototype(resultType, shape(a.shape().d0(), u(), a.shape().d2(), a.shape().d3())))
                    .operation(ptr -> method.apply(ptr, a.dereference(), dim.index()));
 
     }
 
-    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Operation.Builder.Unary<IT, D0, D1, D2, D3>.Single<T, D0, D1, U, D3> reduce(
-        String name, Tensor<IT, D0, D1, D2, D3> a,
+    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Operation.Builder.Unary<Tensor<IT, Shape<D0, D1, D2, D3>>>.Single<Tensor<T, Shape<D0, D1, U, D3>>> reduce(
+        String name, Tensor<IT, Shape<D0, D1, D2, D3>> a,
         Functions.Function3<MemorySegment, MemorySegment, Integer, Integer> method, arrayfire.D2 dim, T resultType) {
         return operation(name)
                    .inputs(a)
-                   .outputs(prototype(resultType, shape(a.d0(), a.d1(), u(), a.d3())))
+                   .outputs(prototype(resultType, shape(a.shape().d0(), a.shape().d1(), u(), a.shape().d3())))
                    .operation(ptr -> method.apply(ptr, a.dereference(), dim.index()));
     }
 
-    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Operation.Builder.Unary<IT, D0, D1, D2, D3>.Single<T, D0, D1, D2, U> reduce(
-        String name, Tensor<IT, D0, D1, D2, D3> a,
+    private static <T extends DataType<?, ?>, IT extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Operation.Builder.Unary<Tensor<IT, Shape<D0, D1, D2, D3>>>.Single<Tensor<T, Shape<D0, D1, D2, U>>> reduce(
+        String name, Tensor<IT, Shape<D0, D1, D2, D3>> a,
         Functions.Function3<MemorySegment, MemorySegment, Integer, Integer> method, arrayfire.D3 dim, T resultType) {
         return operation(name)
                    .inputs(a)
-                   .outputs(prototype(resultType, shape(a.d0(), a.d1(), a.d2(), u())))
+                   .outputs(prototype(resultType, shape(a.shape().d0(), a.shape().d1(), a.shape().d2(), u())))
                    .operation(ptr -> method.apply(ptr, a.dereference(), dim.index()));
     }
 
@@ -428,48 +427,44 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
      * Cast the given tensor to the given type.
      */
     @SuppressWarnings("unchecked")
-    public static <T extends DataType<?, ?>, OT extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<OT, D0, D1, D2, D3> cast(
-        Tensor<T, D0, D1, D2, D3> input, OT type) {
+    public static <T extends DataType<?, ?>, OT extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<OT, S> cast(
+        Tensor<T, S> input, OT type) {
         if (input.type().equals(type)) {
-            return (Tensor<OT, D0, D1, D2, D3>) input;
+            return (Tensor<OT, S>) input;
         }
         return operation("cast")
                    .inputs(input)
                    .outputs(prototype(type, input.shape()))
                    .operation(ptr -> arrayfire_h.af_cast(ptr, input.dereference(), type.code()))
-                   .grads((result, grads) -> grads.cast(input.type()))
+                   .grads((result, grads) -> cast(grads, input.type()))
                    .build();
     }
 
     /**
      * Returns a tensor of value 1 with the same type and shape as the given tensor.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> ones(
-        Tensor<T, D0, D1, D2, D3> model) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> ones(Tensor<T, S> model) {
         return ones(model.type(), model.shape());
     }
 
     /**
      * Returns a tensor of value 1 with the given type and shape.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> ones(
-        T type, Shape<D0, D1, D2, D3> shape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> ones(T type, S shape) {
         return constant(type, 1).tileAs(shape);
     }
 
     /**
      * Returns a tensor of value 0 with the given type and shape.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> zeros(
-        T type, Shape<D0, D1, D2, D3> shape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> zeros(T type, S shape) {
         return constant(type, 0).tileAs(shape);
     }
 
     /**
      * Create a random tensor sampled from uniform distribution between [0, 1].
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> randu(
-        T type, Shape<D0, D1, D2, D3> shape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> randu(T type, S shape) {
         return operation("randu")
                    .inputs()
                    .outputs(prototype(type, shape))
@@ -480,8 +475,7 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
     /**
      * Create a random tensor sampled from a normal distribution with mean 0.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> randn(
-        T type, Shape<D0, D1, D2, D3> shape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> randn(T type, S shape) {
         return operation("randn")
                    .inputs()
                    .outputs(prototype(type, shape))
@@ -492,14 +486,14 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
     /**
      * Create a tensor with values [0, n-1].
      */
-    public static Tensor<U32, N, U, U, U> range(int n) {
+    public static Tensor<U32, Shape<N, U, U, U>> range(int n) {
         return range(U32, n);
     }
 
     /**
      * Create a tensor with values [0, n-1] of the given type.
      */
-    public static <T extends DataType<?, ?>> Tensor<T, N, U, U, U> range(T type, int n) {
+    public static <T extends DataType<?, ?>> Tensor<T, Shape<N, U, U, U>> range(T type, int n) {
         var shape = shape(n(n));
         return operation("range")
                    .inputs()
@@ -525,13 +519,13 @@ public static void setRandomEngineType(RandomEngineType type) {
     /**
      * Pull data from the device to the host, returning a native array.
      */
-    public static <AT extends NativeArray<?, ?, ?>, T extends DataType<AT, ?>> AT data(Tensor<T, ?, ?, ?, ?> a) {
+    public static <AT extends NativeArray<?, ?, ?>, T extends DataType<AT, ?>> AT data(Tensor<T, ?> a) {
         var result = a.type().create(a.capacity());
         handleStatus(() -> arrayfire_h.af_get_data_ptr(result.segment(), a.dereference()));
         return result;
     }
 
-    private static void checkDims(Tensor<?, ?, ?, ?, ?> tensor) {
+    private static void checkDims(Tensor<?, ?> tensor) {
         try (Arena arena = Arena.ofConfined()) {
             var dims = arena.allocateArray(ValueLayout.JAVA_LONG, 4);
             handleStatus(
@@ -641,11 +635,12 @@ private static MemorySegment nativeDims(Shape<?, ?, ?, ?> shape) {
     /**
      * Transpose D0 and D1 dimensions of the given tensor.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D1, D0, D2, D3> transpose(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D1, D0, D2, D3>> transpose(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return operation("transpose")
                    .inputs(tensor)
-                   .outputs(prototype(tensor.type(), shape(tensor.d1(), tensor.d0(), tensor.d2(), tensor.d3())))
+                   .outputs(prototype(tensor.type(),
+                       shape(tensor.shape().d1(), tensor.shape().d0(), tensor.shape().d2(), tensor.shape().d3())))
                    .operation(ptr -> arrayfire_h.af_transpose(ptr, tensor.dereference(), true))
                    .grads((result, grads) -> transpose(grads))
                    .build();
@@ -654,47 +649,47 @@ private static MemorySegment nativeDims(Shape<?, ?, ?, ?> shape) {
     /**
      * Change the type of the tensor's D0 dimension to the given type variable provider.
      */
-    public static <T extends DataType<?, ?>, OD0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, OD0, D1, D2, D3> castshape(
-        Tensor<T, ?, D1, D2, D3> tensor, Function<Integer, OD0> d0) {
-        return reshape(tensor, shape(d0.apply(tensor.d0().size()), tensor.d1(), tensor.d2(), tensor.d3()));
+    public static <T extends DataType<?, ?>, OD0 extends Num<OD0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<OD0, D1, D2, D3>> castshape(
+        Tensor<T, Shape<?, D1, D2, D3>> tensor, Function<Integer, OD0> d0) {
+        return reshape(tensor,
+            shape(d0.apply(tensor.shape().d0().size()), tensor.shape().d1(), tensor.shape().d2(), tensor.shape().d3()));
     }
 
     /**
      * Change the type of the tensor's D0, D1 dimensions to the given type variable providers.
      */
-    public static <T extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, OD0, OD1, D2, D3> castshape(
-        Tensor<T, ?, ?, D2, D3> tensor, Function<Integer, OD0> d0, Function<Integer, OD1> d1) {
+    public static <T extends DataType<?, ?>, OD0 extends Num<OD0>, OD1 extends Num<OD1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<OD0, OD1, D2, D3>> castshape(
+        Tensor<T, Shape<?, ?, D2, D3>> tensor, Function<Integer, OD0> d0, Function<Integer, OD1> d1) {
         return reshape(tensor,
-            shape(d0.apply(tensor.d0().size()), d1.apply(tensor.d1().size()), tensor.d2(), tensor.d3()));
+            shape(d0.apply(tensor.shape().d0().size()), d1.apply(tensor.shape().d1().size()), tensor.shape().d2(),
+                tensor.shape().d3()));
     }
 
     /**
      * Change the type of the tensor's D0, D1, D2 dimensions to the given type variable providers.
      */
-    public static <T extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, D3 extends Num<?>> Tensor<T, OD0, OD1, OD2, D3> castshape(
-        Tensor<T, ?, ?, ?, D3> tensor, Function<Integer, OD0> d0, Function<Integer, OD1> d1,
+    public static <T extends DataType<?, ?>, OD0 extends Num<OD0>, OD1 extends Num<OD1>, OD2 extends Num<OD2>, D3 extends Num<D3>> Tensor<T, Shape<OD0, OD1, OD2, D3>> castshape(
+        Tensor<T, Shape<?, ?, ?, D3>> tensor, Function<Integer, OD0> d0, Function<Integer, OD1> d1,
         Function<Integer, OD2> d2) {
-        return reshape(tensor,
-            shape(d0.apply(tensor.d0().size()), d1.apply(tensor.d1().size()), d2.apply(tensor.d2().size()),
-                tensor.d3()));
+        return reshape(tensor, shape(d0.apply(tensor.shape().d0().size()), d1.apply(tensor.shape().d1().size()),
+            d2.apply(tensor.shape().d2().size()), tensor.shape().d3()));
     }
 
     /**
      * Change the type of the tensor's dimensions to the given type variable providers.
      */
-    public static <T extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> castshape(
-        Tensor<T, ?, ?, ?, ?> tensor, Function<Integer, OD0> d0, Function<Integer, OD1> d1, Function<Integer, OD2> d2,
-        Function<Integer, OD3> d3) {
-        return reshape(tensor,
-            shape(d0.apply(tensor.d0().size()), d1.apply(tensor.d1().size()), d2.apply(tensor.d2().size()),
-                d3.apply(tensor.d3().size())));
+    public static <T extends DataType<?, ?>, OD0 extends Num<OD0>, OD1 extends Num<OD1>, OD2 extends Num<OD2>, OD3 extends Num<OD3>> Tensor<T, Shape<OD0, OD1, OD2, OD3>> castshape(
+        Tensor<T, Shape<?, ?, ?, ?>> tensor, Function<Integer, OD0> d0, Function<Integer, OD1> d1,
+        Function<Integer, OD2> d2, Function<Integer, OD3> d3) {
+        return reshape(tensor, shape(d0.apply(tensor.shape().d0().size()), d1.apply(tensor.shape().d1().size()),
+            d2.apply(tensor.shape().d2().size()), d3.apply(tensor.shape().d3().size())));
     }
 
     /**
      * Reshape the tensor to the given shape.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> reshape(
-        Tensor<T, D0, D1, D2, D3> tensor, Shape<OD0, OD1, OD2, OD3> newShape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>, NS extends Shape<?, ?, ?, ?>> Tensor<T, NS> reshape(
+        Tensor<T, S> tensor, NS newShape) {
         if (tensor.shape().capacity() != newShape.capacity()) {
             throw new IllegalArgumentException(
                 String.format("New shape %s doesn't have same capacity as original shape %s", newShape,
@@ -712,7 +707,7 @@ private static MemorySegment nativeDims(Shape<?, ?, ?, ?> shape) {
     /**
      * Release the memory of the given tensor on the device.
      */
-    public static void release(Tensor<?, ?, ?, ?, ?> tensor) {
+    public static void release(Tensor<?, ?> tensor) {
         handleStatus(() -> arrayfire_h.af_release_array(tensor.dereference()));
         Scope.untrack(tensor);
     }
@@ -720,8 +715,7 @@ public static void release(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Retain the given tensor, increasing its ref count by 1 and return a new container for it.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> retain(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <DT extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<DT, S> retain(Tensor<DT, S> tensor) {
         return operation("retain")
                    .inputs(tensor)
                    .outputs(prototype(tensor.type(), tensor.shape()))
@@ -733,8 +727,8 @@ public static void release(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Set the values of the given variable to the values of the given tensor.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Operation set(
-        Variable<T, D0, D1, D2, D3> variable, Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Operation set(Variable<T, S> variable,
+                                                                                        Tensor<T, S> tensor) {
         return operation("set").inputs(tensor).outputs().operation(() -> {
             handleStatus(() -> arrayfire_h.af_release_array(variable.dereference()));
             handleStatus(() -> arrayfire_h.af_retain_array(variable.segment(), tensor.dereference()));
@@ -744,7 +738,7 @@ public static void release(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Return the ref count of the given tensor.
      */
-    public static int refCount(Tensor<?, ?, ?, ?, ?> tensor) {
+    public static int refCount(Tensor<?, ?> tensor) {
         try (Arena arena = Arena.ofConfined()) {
             var result = arena.allocate(ValueLayout.JAVA_INT);
             handleStatus(() -> arrayfire_h.af_get_data_ref_count(result, tensor.dereference()));
@@ -755,8 +749,8 @@ public static int refCount(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Create a variable with the given initializer.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Variable<T, D0, D1, D2, D3> variable(
-        Supplier<Tensor<T, D0, D1, D2, D3>> initializer) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Variable<T, S> variable(
+        Supplier<Tensor<T, S>> initializer) {
         var tensor = af.tidy(initializer);
         var variable = new Variable<>(tensor.type(), tensor.shape());
         variable.segment().copyFrom(tensor.segment());
@@ -767,8 +761,8 @@ public static int refCount(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Create params with the given initializer and optimizer.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Params<T, D0, D1, D2, D3> params(
-        Supplier<Tensor<T, D0, D1, D2, D3>> initializer, OptimizerProvider optimizerProvider) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Params<T, S> params(
+        Supplier<Tensor<T, S>> initializer, OptimizerProvider optimizerProvider) {
         var tensor = af.tidy(initializer);
         var params = new Params<>(tensor.type(), tensor.shape(), optimizerProvider);
         params.segment().copyFrom(tensor.segment());
@@ -779,8 +773,7 @@ public static int refCount(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Evaluate the tensor, telling the ArrayFire JIT compiler that you want the literal values of the tensor.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> eval(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> eval(Tensor<T, S> tensor) {
         handleStatus(() -> arrayfire_h.af_eval(tensor.dereference()));
         return tensor;
     }
@@ -788,7 +781,7 @@ public static int refCount(Tensor<?, ?, ?, ?, ?> tensor) {
     /**
      * Evaluate the tensors, telling the ArrayFire JIT compiler that you want the literal values of the tensors.
      */
-    public static void eval(Tensor<?, ?, ?, ?, ?>... tensors) {
+    public static void eval(Tensor<?, ?>... tensors) {
         try (Arena arena = Arena.ofConfined()) {
             var array = arena.allocateArray(ValueLayout.ADDRESS, tensors.length);
             for (int i = 0; i < tensors.length; i++) {
@@ -801,15 +794,17 @@ public static void eval(Tensor<?, ?, ?, ?, ?>... tensors) {
     /**
      * Multiply two tensors together element wise, broadcasting the smaller tensor to the larger tensor's shape.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> mul(
-        Tensor<T, D0, D1, D2, D3> tensor, Tileable<T, ?, ?, ?, ?> tileable) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> mul(Tensor<T, S> tensor,
+                                                                                           Tileable<T, ?> tileable) {
         checkTileableIsSmaller(tensor, tileable);
         return mul(tensor, tileable.tensor().tileAs(tensor));
     }
 
-    private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<?, ?, ?, ?, ?> right) {
-        if (left.d0().size() < right.tensor().d0().size() || left.d1().size() < right.tensor().d1().size() ||
-                left.d2().size() < right.tensor().d2().size() || left.d3().size() < right.tensor().d3().size()) {
+    private static void checkTileableIsSmaller(Tensor<?, ?> left, Tileable<?, ?> right) {
+        if (left.shape().d0().size() < right.tensor().shape().d0().size() ||
+                left.shape().d1().size() < right.tensor().shape().d1().size() ||
+                left.shape().d2().size() < right.tensor().shape().d2().size() ||
+                left.shape().d3().size() < right.tensor().shape().d3().size()) {
             throw new IllegalArgumentException(
                 String.format("Tileable shape %s is larger than tensor shape %s", right.tensor().shape(),
                     left.shape()));
@@ -819,16 +814,16 @@ private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<
     /**
      * Multiply the tensor by a scalar value.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> mul(
-        Tensor<T, D0, D1, D2, D3> left, double right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> mul(Tensor<T, S> left,
+                                                                                           double right) {
         return mul(left, af.constant(left.type(), left.shape(), right));
     }
 
     /**
      * Multiply two tensors together, element wise.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> mul(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> mul(Tensor<T, S> left,
+                                                                                           Tensor<T, S> right) {
         return operation("mul")
                    .inputs(left, right)
                    .outputs(prototype(left.type(), left.shape()))
@@ -837,24 +832,24 @@ private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> div(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> div(Tensor<T, S> left,
+                                                                                           Tensor<T, S> right) {
         return operation("div")
                    .inputs(left, right)
                    .outputs(prototype(left.type(), left.shape()))
                    .operation(ptr -> arrayfire_h.af_div(ptr, left.dereference(), right.dereference(), true))
                    .grads((result, grads) -> {
-                       var rightReciprocal = af.div(af.constant(1f).cast(left.type()).tileAs(right), right);
+                       var rightReciprocal = div(constant(1f).cast(left.type()).tileAs(right), right);
                        var leftGrads = mul(rightReciprocal, grads);
-                       var rightGrads = af.mul(af.mul(leftGrads, left.negate()), rightReciprocal);
+                       var rightGrads = mul(mul(leftGrads, left.negate()), rightReciprocal);
                        return new TensorPair<>(leftGrads, rightGrads);
                    })
                    .build();
 
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> add(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> add(Tensor<T, S> left,
+                                                                                           Tensor<T, S> right) {
         return operation("add")
                    .inputs(left, right)
                    .outputs(prototype(left.type(), left.shape()))
@@ -863,8 +858,8 @@ private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> sub(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sub(Tensor<T, S> left,
+                                                                                           Tensor<T, S> right) {
         return operation("sub")
                    .inputs(left, right)
                    .outputs(prototype(left.type(), left.shape()))
@@ -873,8 +868,8 @@ private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<B8, D0, D1, D2, D3> ge(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, D0, D1, D2, D3> rhs) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<B8, S> ge(Tensor<T, S> tensor,
+                                                                                           Tensor<T, S> rhs) {
         return operation("ge")
                    .inputs(tensor, rhs)
                    .outputs(prototype(B8, tensor.shape()))
@@ -882,8 +877,8 @@ private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<B8, D0, D1, D2, D3> le(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, D0, D1, D2, D3> rhs) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<B8, S> le(Tensor<T, S> tensor,
+                                                                                           Tensor<T, S> rhs) {
         return operation("le")
                    .inputs(tensor, rhs)
                    .outputs(prototype(B8, tensor.shape()))
@@ -891,8 +886,7 @@ private static void checkTileableIsSmaller(Tensor<?, ?, ?, ?, ?> left, Tileable<
                    .build();
     }
 
-    public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<B8, D0, D1, D2, D3> and(
-        Tensor<B8, D0, D1, D2, D3> left, Tensor<B8, D0, D1, D2, D3> right) {
+    public static <S extends Shape<?, ?, ?, ?>> Tensor<B8, S> and(Tensor<B8, S> left, Tensor<B8, S> right) {
         return operation("and")
                    .inputs(left, right)
                    .outputs(prototype(B8, left.shape()))
@@ -900,271 +894,277 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> maxof(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> maxof(Tensor<T, S> left,
+                                                                                             Tensor<T, S> right) {
         return operation("maxof")
                    .inputs(left, right)
                    .outputs(prototype(left.type(), left.shape()))
                    .operation(ptr -> arrayfire_h.af_maxof(ptr, left.dereference(), right.dereference(), true))
                    .grads((result, grads) -> {
-                       var leftIsMax = af.eq(result, left).cast(left.type());
-                       var rightIsMax = af.eq(result, right).cast(left.type());
+                       var leftIsMax = eq(result, left).cast(left.type());
+                       var rightIsMax = eq(result, right).cast(left.type());
                        return new TensorPair<>(mul(leftIsMax, grads), mul(rightIsMax, grads));
                    })
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> minof(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> minof(Tensor<T, S> left,
+                                                                                             Tensor<T, S> right) {
         return operation("minof")
                    .inputs(left, right)
                    .outputs(prototype(left.type(), left.shape()))
                    .operation(ptr -> arrayfire_h.af_minof(ptr, left.dereference(), right.dereference(), true))
                    .grads((result, grads) -> {
-                       var leftIsMin = af.eq(result, left).cast(left.type());
-                       var rightIsMin = af.eq(result, right).cast(left.type());
+                       var leftIsMin = eq(result, left).cast(left.type());
+                       var rightIsMin = eq(result, right).cast(left.type());
                        return new TensorPair<>(mul(leftIsMin, grads), mul(rightIsMin, grads));
                    })
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, LD0 extends Num<LD0>, RD0 extends Num<RD0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, N, D1, D2, D3> join(
-        Tensor<T, LD0, D1, D2, D3> lhs, Tensor<T, RD0, D1, D2, D3> rhs) {
+    public static <T extends DataType<?, ?>, LD0 extends Num<LD0>, RD0 extends Num<RD0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<N, D1, D2, D3>> join(
+        Tensor<T, Shape<LD0, D1, D2, D3>> lhs, Tensor<T, Shape<RD0, D1, D2, D3>> rhs) {
         return operation("join")
                    .inputs(lhs, rhs)
-                   .outputs(
-                       prototype(lhs.type(), shape(n(lhs.d0().size() + rhs.d0().size()), lhs.d1(), lhs.d2(), lhs.d3())))
+                   .outputs(prototype(lhs.type(),
+                       shape(n(lhs.shape().d0().size() + rhs.shape().d0().size()), lhs.shape().d1(), lhs.shape().d2(),
+                           lhs.shape().d3())))
                    .operation(ptr -> arrayfire_h.af_join(ptr, 0, lhs.dereference(), rhs.dereference()))
-                   .grads((result, grads) -> new TensorPair<>(index(grads, seq(lhs.d0())),
-                       index(grads, seq(lhs.d0().size(), rhs.d0()))))
+                   .grads((result, grads) -> new TensorPair<>(index(grads, seq(lhs.shape().d0())),
+                       index(grads, seq(lhs.shape().d0().size(), rhs.shape().d0()))))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, LD1 extends Num<LD1>, RD1 extends Num<RD1>, D0 extends Num<D0>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, N, D2, D3> join(
-        Tensor<T, D0, LD1, D2, D3> lhs, Tensor<T, D0, RD1, D2, D3> rhs, arrayfire.D1 ignored) {
-        if (!(lhs.d0().size() == rhs.d0().size() && lhs.d2().size() == rhs.d2().size() &&
-                  lhs.d3().size() == rhs.d3().size())) {
+    public static <T extends DataType<?, ?>, LD1 extends Num<LD1>, RD1 extends Num<RD1>, D0 extends Num<D0>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, N, D2, D3>> join(
+        Tensor<T, Shape<D0, LD1, D2, D3>> lhs, Tensor<T, Shape<D0, RD1, D2, D3>> rhs, arrayfire.D1 ignored) {
+        if (!(lhs.shape().d0().size() == rhs.shape().d0().size() &&
+                  lhs.shape().d2().size() == rhs.shape().d2().size() &&
+                  lhs.shape().d3().size() == rhs.shape().d3().size())) {
             throw new IllegalArgumentException(
                 String.format("Incompatible shapes to join along d1: %s, %s", lhs.shape(), rhs.shape()));
         }
         return operation("join")
                    .inputs(lhs, rhs)
-                   .outputs(
-                       prototype(lhs.type(), shape(lhs.d0(), n(lhs.d1().size() + rhs.d1().size()), lhs.d2(), lhs.d3())))
+                   .outputs(prototype(lhs.type(),
+                       shape(lhs.shape().d0(), n(lhs.shape().d1().size() + rhs.shape().d1().size()), lhs.shape().d2(),
+                           lhs.shape().d3())))
                    .operation(ptr -> arrayfire_h.af_join(ptr, 1, lhs.dereference(), rhs.dereference()))
-                   .grads((result, grads) -> new TensorPair<>(index(grads, span(), seq(lhs.d1())),
-                       index(grads, span(), seq(lhs.d1().size(), rhs.d1()))))
+                   .grads((result, grads) -> new TensorPair<>(index(grads, span(), seq(lhs.shape().d1())),
+                       index(grads, span(), seq(lhs.shape().d1().size(), rhs.shape().d1()))))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, LD2 extends Num<LD2>, RD2 extends Num<RD2>, D0 extends Num<D0>, D1 extends Num<D1>, D3 extends Num<D3>> Tensor<T, D0, D1, N, D3> join(
-        Tensor<T, D0, D1, LD2, D3> lhs, Tensor<T, D0, D1, RD2, D3> rhs, arrayfire.D2 ignored) {
-        if (!(lhs.d0().size() == rhs.d0().size() && lhs.d1().size() == rhs.d1().size() &&
-                  lhs.d3().size() == rhs.d3().size())) {
+    public static <T extends DataType<?, ?>, LD2 extends Num<LD2>, RD2 extends Num<RD2>, D0 extends Num<D0>, D1 extends Num<D1>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, N, D3>> join(
+        Tensor<T, Shape<D0, D1, LD2, D3>> lhs, Tensor<T, Shape<D0, D1, RD2, D3>> rhs, arrayfire.D2 ignored) {
+        if (!(lhs.shape().d0().size() == rhs.shape().d0().size() &&
+                  lhs.shape().d1().size() == rhs.shape().d1().size() &&
+                  lhs.shape().d3().size() == rhs.shape().d3().size())) {
             throw new IllegalArgumentException(
                 String.format("Incompatible shapes to join along d2: %s, %s", lhs.shape(), rhs.shape()));
         }
         return operation("join")
                    .inputs(lhs, rhs)
-                   .outputs(
-                       prototype(lhs.type(), shape(lhs.d0(), lhs.d1(), n(lhs.d2().size() + rhs.d2().size()), lhs.d3())))
+                   .outputs(prototype(lhs.type(),
+                       shape(lhs.shape().d0(), lhs.shape().d1(), n(lhs.shape().d2().size() + rhs.shape().d2().size()),
+                           lhs.shape().d3())))
                    .operation(ptr -> arrayfire_h.af_join(ptr, 2, lhs.dereference(), rhs.dereference()))
-                   .grads((result, grads) -> new TensorPair<>(index(grads, span(), span(), seq(lhs.d2())),
-                       index(grads, span(), span(), seq(lhs.d2().size(), rhs.d2()))))
+                   .grads((result, grads) -> new TensorPair<>(index(grads, span(), span(), seq(lhs.shape().d2())),
+                       index(grads, span(), span(), seq(lhs.shape().d2().size(), rhs.shape().d2()))))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, LD3 extends Num<LD3>, RD3 extends Num<RD3>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>> Tensor<T, D0, D1, D2, N> join(
-        Tensor<T, D0, D1, D2, LD3> lhs, Tensor<T, D0, D1, D2, RD3> rhs, arrayfire.D3 ignored) {
-        if (!(lhs.d0().size() == rhs.d0().size() && lhs.d1().size() == rhs.d1().size() &&
-                  lhs.d2().size() == rhs.d2().size())) {
+    public static <T extends DataType<?, ?>, LD3 extends Num<LD3>, RD3 extends Num<RD3>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>> Tensor<T, Shape<D0, D1, D2, N>> join(
+        Tensor<T, Shape<D0, D1, D2, LD3>> lhs, Tensor<T, Shape<D0, D1, D2, RD3>> rhs, arrayfire.D3 ignored) {
+        if (!(lhs.shape().d0().size() == rhs.shape().d0().size() &&
+                  lhs.shape().d1().size() == rhs.shape().d1().size() &&
+                  lhs.shape().d2().size() == rhs.shape().d2().size())) {
             throw new IllegalArgumentException(
                 String.format("Incompatible shapes to join along d3: %s, %s", lhs.shape(), rhs.shape()));
         }
         return operation("join")
                    .inputs(lhs, rhs)
-                   .outputs(
-                       prototype(lhs.type(), shape(lhs.d0(), lhs.d1(), lhs.d2(), n(lhs.d3().size() + rhs.d3().size()))))
+                   .outputs(prototype(lhs.type(), shape(lhs.shape().d0(), lhs.shape().d1(), lhs.shape().d2(),
+                       n(lhs.shape().d3().size() + rhs.shape().d3().size()))))
                    .operation(ptr -> arrayfire_h.af_join(ptr, 3, lhs.dereference(), rhs.dereference()))
-                   .grads((result, grads) -> new TensorPair<>(
-                       index(grads, span(), span(), span(), seq(lhs.d3())),
-                       index(grads, span(), span(), span(), seq(lhs.d3().size(), rhs.d3()))))
+                   .grads(
+                       (result, grads) -> new TensorPair<>(index(grads, span(), span(), span(), seq(lhs.shape().d3())),
+                           index(grads, span(), span(), span(), seq(lhs.shape().d3().size(), rhs.shape().d3()))))
                    .build();
     }
 
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, U, D1, D2, D3> sum(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<U, D1, D2, D3>> sum(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return sum(tensor, D0);
     }
 
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, U, D1, D2, D3> sum(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D0 dim) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<U, D1, D2, D3>> sum(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D0 dim) {
         return reduce("sum", tensor, arrayfire_h::af_sum, dim, tensor.type().sumType())
                    .grads((result, grads) -> grads.cast(tensor.type()).tileAs(tensor))
                    .build();
 
     }
 
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, D0, U, D2, D3> sum(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D1 dim) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<D0, U, D2, D3>> sum(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D1 dim) {
         return reduce("sum", tensor, arrayfire_h::af_sum, dim, tensor.type().sumType())
                    .grads((result, grads) -> grads.cast(tensor.type()).tileAs(tensor))
                    .build();
     }
 
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, D0, D1, U, D3> sum(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D2 dim) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<D0, D1, U, D3>> sum(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D2 dim) {
         return reduce("sum", tensor, arrayfire_h::af_sum, dim, tensor.type().sumType())
                    .grads((result, grads) -> grads.cast(tensor.type()).tileAs(tensor))
                    .build();
     }
 
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, D0, D1, D2, U> sum(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D3 dim) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<D0, D1, D2, U>> sum(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D3 dim) {
         return reduce("sum", tensor, arrayfire_h::af_sum, dim, tensor.type().sumType())
                    .grads((result, grads) -> grads.cast(tensor.type()).tileAs(tensor))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> mean(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> mean(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return mean(tensor, D0);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> mean(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D0 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> mean(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D0 dim) {
         return reduce("mean", tensor, arrayfire_h::af_mean, dim, tensor.type())
                    .grads((result, grads) -> af.div(grads.tileAs(tensor),
-                       af.constant(tensor.type(), tensor.d0().size()).tileAs(tensor)))
+                       af.constant(tensor.type(), tensor.shape().d0().size()).tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, U, D2, D3> mean(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D1 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, U, D2, D3>> mean(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D1 dim) {
         return reduce("mean", tensor, arrayfire_h::af_mean, dim, tensor.type())
                    .grads((result, grads) -> af.div(grads.tileAs(tensor),
-                       af.constant(tensor.type(), tensor.d1().size()).tileAs(tensor)))
+                       af.constant(tensor.type(), tensor.shape().d1().size()).tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, U, D3> mean(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D2 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, U, D3>> mean(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D2 dim) {
         return reduce("mean", tensor, arrayfire_h::af_mean, dim, tensor.type())
                    .grads((result, grads) -> af.div(grads.tileAs(tensor),
-                       af.constant(tensor.type(), tensor.d2().size()).tileAs(tensor)))
+                       af.constant(tensor.type(), tensor.shape().d2().size()).tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, U> mean(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D3 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, U>> mean(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D3 dim) {
         return reduce("mean", tensor, arrayfire_h::af_mean, dim, tensor.type())
                    .grads((result, grads) -> af.div(grads.tileAs(tensor),
-                       af.constant(tensor.type(), tensor.d3().size()).tileAs(tensor)))
+                       af.constant(tensor.type(), tensor.shape().d3().size()).tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> median(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> median(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return median(tensor, D0);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> median(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D0 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> median(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D0 dim) {
         return reduce("median", tensor, arrayfire_h::af_median, dim, tensor.type()).build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, U, D2, D3> median(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D1 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, U, D2, D3>> median(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D1 dim) {
         return reduce("median", tensor, arrayfire_h::af_median, dim, tensor.type()).build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, U, D3> median(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D2 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, U, D3>> median(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D2 dim) {
         return reduce("median", tensor, arrayfire_h::af_median, dim, tensor.type()).build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, U> median(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D3 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, U>> median(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D3 dim) {
         return reduce("median", tensor, arrayfire_h::af_median, dim, tensor.type()).build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> max(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> max(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return max(tensor, D0);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> max(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D0 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> max(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D0 dim) {
         return reduce("max", tensor, arrayfire_h::af_max, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, U, D2, D3> max(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D1 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, U, D2, D3>> max(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D1 dim) {
         return reduce("max", tensor, arrayfire_h::af_max, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, U, D3> max(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D2 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, U, D3>> max(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D2 dim) {
         return reduce("max", tensor, arrayfire_h::af_max, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, U> max(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D3 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, U>> max(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D3 dim) {
         return reduce("max", tensor, arrayfire_h::af_max, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> min(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> min(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return min(tensor, D0);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, U, D1, D2, D3> min(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D0 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<U, D1, D2, D3>> min(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D0 dim) {
         return reduce("min", tensor, arrayfire_h::af_min, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, U, D2, D3> min(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D1 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, U, D2, D3>> min(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D1 dim) {
         return reduce("min", tensor, arrayfire_h::af_min, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, U, D3> min(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D2 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, U, D3>> min(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D2 dim) {
         return reduce("min", tensor, arrayfire_h::af_min, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, U> min(
-        Tensor<T, D0, D1, D2, D3> tensor, arrayfire.D3 dim) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, U>> min(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, arrayfire.D3 dim) {
         return reduce("min", tensor, arrayfire_h::af_min, dim, tensor.type())
                    .grads((result, grads) -> mul(af.eq(result.tileAs(tensor), tensor).cast(grads.type()),
                        grads.tileAs(tensor)))
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> ImaxResult<T, U, D1, D2, D3> imax(
-        Tensor<T, D0, D1, D2, D3> tensor) {
-        var shape = shape(u(), tensor.d1(), tensor.d2(), tensor.d3());
+    public static <T extends DataType<?, ?>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> ImaxResult<T, Shape<U, D1, D2, D3>> imax(
+        Tensor<T, ? extends Shape<?, D1, D2, D3>> tensor) {
+        var shape = shape(u(), tensor.shape().d1(), tensor.shape().d2(), tensor.shape().d3());
         var pair = operation("imax")
                        .inputs(tensor)
                        .outputs(prototype(tensor.type(), shape), prototype(U32, shape))
@@ -1174,9 +1174,9 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
         return new ImaxResult<>(pair.left(), pair.right());
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, K extends Num<?>> TopKResult<T, K, D1, D2, D3> topk(
-        Tensor<T, D0, D1, D2, D3> tensor, K k) {
-        var shape = shape(k, tensor.d1(), tensor.d2(), tensor.d3());
+    public static <T extends DataType<?, ?>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, K extends Num<K>> TopKResult<T, Shape<K, D1, D2, D3>> topk(
+        Tensor<T, Shape<?, D1, D2, D3>> tensor, K k) {
+        var shape = shape(k, tensor.shape().d1(), tensor.shape().d2(), tensor.shape().d3());
         var pair = operation("topk")
                        .inputs(tensor)
                        .outputs(prototype(tensor.type(), shape), prototype(U32, shape))
@@ -1187,38 +1187,41 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
         return new TopKResult<>(pair.left(), pair.right());
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D0, D2, D3> diag(
-        Tensor<T, D0, U, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D0, D2, D3>> diag(
+        Tensor<T, Shape<D0, U, D2, D3>> tensor) {
         return operation("diag")
                    .inputs(tensor)
-                   .outputs(prototype(tensor.type(), shape(tensor.d0(), tensor.d0(), tensor.d2(), tensor.d3())))
+                   .outputs(prototype(tensor.type(),
+                       shape(tensor.shape().d0(), tensor.shape().d0(), tensor.shape().d2(), tensor.shape().d3())))
                    .operation(ptr -> arrayfire_h.af_diag_create(ptr, tensor.dereference(), 0))
                    // TODO: Implement grad function.
                    .build();
     }
 
     // https://arrayfire.org/docs/group__blas__func__matmul.htm
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, OD1 extends Num<?>> Tensor<T, D0, OD1, D2, D3> matmul(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D1, OD1, D2, D3> right) {
-        if (left.d1().size() != right.d0().size()) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, OD1 extends Num<OD1>> Tensor<T, Shape<D0, OD1, D2, D3>> matmul(
+        Tensor<T, Shape<D0, D1, D2, D3>> left, Tensor<T, Shape<D1, OD1, D2, D3>> right) {
+        if (left.shape().d1().size() != right.shape().d0().size()) {
             throw new IllegalArgumentException(
                 String.format("Incompatible shapes for matmul, left: %s right: %s", left.shape(), right.shape()));
         }
         return operation("matmul")
                    .inputs(left, right)
-                   .outputs(prototype(left.type(), shape(left.d0(), right.d1(), left.d2(), left.d3())))
+                   .outputs(prototype(left.type(),
+                       shape(left.shape().d0(), right.shape().d1(), left.shape().d2(), left.shape().d3())))
                    .operation(ptr -> arrayfire_h.af_matmul(ptr, left.dereference(), right.dereference(), 0, 0))
                    .grads((result, grads) -> {
-                       var leftGrads = matmul(grads, right.transpose());
-                       var rightGrads = matmul(left.transpose(), grads);
+                       var leftGrads = matmul(grads, transpose(right));
+                       var rightGrads = matmul(transpose(left), grads);
                        return new TensorPair<>(leftGrads, rightGrads);
                    })
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, AD0 extends Num<?>, AD1 extends Num<?>, BD1 extends Num<?>, CD1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, AD0, CD1, D2, D3> matmul(
-        Tensor<T, AD0, AD1, D2, D3> a, Tensor<T, AD1, BD1, D2, D3> b, Tensor<T, BD1, CD1, D2, D3> c) {
-        if (a.d0().size() * b.d1().size() < b.d0().size() * c.d1().size()) {
+    public static <T extends DataType<?, ?>, AD0 extends Num<AD0>, AD1 extends Num<AD1>, BD1 extends Num<BD1>, CD1 extends Num<CD1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<AD0, CD1, D2, D3>> matmul(
+        Tensor<T, Shape<AD0, AD1, D2, D3>> a, Tensor<T, Shape<AD1, BD1, D2, D3>> b,
+        Tensor<T, Shape<BD1, CD1, D2, D3>> c) {
+        if (a.shape().d0().size() * b.shape().d1().size() < b.shape().d0().size() * c.shape().d1().size()) {
             var tmp = matmul(a, b);
             var result = matmul(tmp, c);
             tmp.release();
@@ -1231,8 +1234,9 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
         }
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> clamp(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, D0, D1, D2, D3> lo, Tensor<T, D0, D1, D2, D3> hi) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> clamp(Tensor<T, S> tensor,
+                                                                                             Tensor<T, S> lo,
+                                                                                             Tensor<T, S> hi) {
         return operation("clamp")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1247,14 +1251,13 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
 
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> relu(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> relu(Tensor<T, S> tensor) {
         return clamp(tensor, constant(tensor.type(), 0f).tileAs(tensor),
             constant(tensor.type(), Double.POSITIVE_INFINITY).tileAs(tensor));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<B8, D0, D1, D2, D3> eq(
-        Tensor<T, D0, D1, D2, D3> left, Tensor<T, D0, D1, D2, D3> right) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<B8, S> eq(Tensor<T, S> left,
+                                                                                           Tensor<T, S> right) {
         return operation("eq")
                    .inputs(left, right)
                    .outputs(prototype(B8, left.shape()))
@@ -1262,14 +1265,12 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> negate(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> negate(Tensor<T, S> tensor) {
         var minusOne = constant(tensor.type(), tensor.shape(), -1);
         return mul(tensor, minusOne);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> exp(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> exp(Tensor<T, S> tensor) {
         return operation("exp")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1278,13 +1279,13 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> pow(
-        Tensor<T, D0, D1, D2, D3> tensor, double pow) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> pow(Tensor<T, S> tensor,
+                                                                                           double pow) {
         return pow(tensor, constant(tensor.type(), tensor.shape(), pow));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> pow(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, D0, D1, D2, D3> pow) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> pow(Tensor<T, S> tensor,
+                                                                                           Tensor<T, S> pow) {
         return operation("pow")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1297,8 +1298,7 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
     /**
      * Returns 1 for negative numbers and 0 for positive numbers.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> signbit(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> signbit(Tensor<T, S> tensor) {
         return operation("signbit")
                    .inputs(tensor)
                    .outputs(tensor.prototype())
@@ -1309,8 +1309,7 @@ public static <D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 exten
     /**
      * Returns -1 for negative numbers and 1 for positive numbers.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> signum(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> signum(Tensor<T, S> tensor) {
         return operation("signum")
                    .inputs(tensor)
                    .outputs(tensor.prototype())
@@ -1323,8 +1322,7 @@ public static Operation.Builder operation(String name) {
         return new Operation.Builder().name(name);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> log(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> log(Tensor<T, S> tensor) {
         return operation("log")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1333,8 +1331,7 @@ public static Operation.Builder operation(String name) {
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> abs(
-        Tensor<T, D0, D1, D2, D3> input) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> abs(Tensor<T, S> input) {
         return operation("abs")
                    .inputs(input)
                    .outputs(prototype(input))
@@ -1343,8 +1340,7 @@ public static Operation.Builder operation(String name) {
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> sqrt(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sqrt(Tensor<T, S> tensor) {
         return operation("sqrt")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1353,12 +1349,12 @@ public static Operation.Builder operation(String name) {
                    .build();
     }
 
-    public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> softmax(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, T>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> softmax(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return softmax(tensor, 1f);
     }
 
-    public static Function<MemorySegment, Integer> tidyOperation(Supplier<Tensor<?, ?, ?, ?, ?>> fn) {
+    public static Function<MemorySegment, Integer> tidyOperation(Supplier<Tensor<?, ?>> fn) {
         return ptr -> {
             var result = tidy(fn);
             ptr.copyFrom(result.segment());
@@ -1367,8 +1363,8 @@ public static Function<MemorySegment, Integer> tidyOperation(Supplier<Tensor<?,
         };
     }
 
-    public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> softmax(
-        Tensor<T, D0, D1, D2, D3> tensor, float temperature) {
+    public static <T extends DataType<?, T>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> softmax(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, float temperature) {
         return operation("softmax").inputs(tensor).outputs(prototype(tensor)).operation(tidyOperation(() -> {
             var max = max(tensor);
             var normalized = sub(tensor, max.tileAs(tensor));
@@ -1378,24 +1374,23 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
             // Compact all dimensions except the first into a batch dimension, so we have a spare dimension for the jacobian.
             var shape = result.shape();
             var workingShape = af.shape(shape.d0(), af.u(),
-                af.b(result.d1().size() * result.d2().size() * result.d3().size()));
+                af.b(result.shape().d1().size() * result.shape().d2().size() * result.shape().d3().size()));
             var resultTensor = result.reshape(workingShape);
             var gradsTensor = grads.reshape(workingShape);
             var positives = af.mul(resultTensor, gradsTensor);
-            var negatives = af.matmul(resultTensor, resultTensor.transpose(), gradsTensor);
+            var negatives = af.matmul(resultTensor, transpose(resultTensor), gradsTensor);
             var inputGrads = af.sub(positives, negatives);
             return inputGrads.reshape(tensor.shape());
         }).build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> sigmoid(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sigmoid(Tensor<T, S> tensor) {
         var one = ones(tensor);
         return div(one, add(one, exp(negate(tensor))));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> sparse(
-        Tensor<T, D0, D1, D2, D3> tensor, Storage storage) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> sparse(Tensor<T, S> tensor,
+                                                                                              Storage storage) {
         return operation("sparse")
                    .inputs(tensor)
                    .outputs(prototype(tensor.type(), tensor.shape()))
@@ -1405,48 +1400,48 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, ?, D1, D2, D3> tensor, Index<D0> i0) {
-        return index(tensor, i0, seq(tensor.d1()), seq(tensor.d2()), seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<?, D1, D2, D3>> tensor, Index<D0> i0) {
+        return index(tensor, i0, seq(tensor.shape().d1()), seq(tensor.shape().d2()), seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, ?, ?, D2, D3> tensor, Index<D0> i0, Index<D1> i1) {
-        return index(tensor, i0, i1, seq(tensor.d2()), seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<?, ?, D2, D3>> tensor, Index<D0> i0, Index<D1> i1) {
+        return index(tensor, i0, i1, seq(tensor.shape().d2()), seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, D0, ?, D2, D3> tensor, Span ignored0, Index<D1> i1) {
-        return index(tensor, seq(tensor.d0()), i1, seq(tensor.d2()), seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<D0, ?, D2, D3>> tensor, Span ignored0, Index<D1> i1) {
+        return index(tensor, seq(tensor.shape().d0()), i1, seq(tensor.shape().d2()), seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, D0, D1, ?, D3> tensor, Span ignored0, Span ignored1, Index<D2> i2) {
-        return index(tensor, seq(tensor.d0()), seq(tensor.d1()), i2, seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<D0, D1, ?, D3>> tensor, Span ignored0, Span ignored1, Index<D2> i2) {
+        return index(tensor, seq(tensor.shape().d0()), seq(tensor.shape().d1()), i2, seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, D0, ?, ?, D3> tensor, Span ignored0, Index<D1> i1, Index<D2> i2) {
-        return index(tensor, seq(tensor.d0()), i1, i2, seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<D0, ?, ?, D3>> tensor, Span ignored0, Index<D1> i1, Index<D2> i2) {
+        return index(tensor, seq(tensor.shape().d0()), i1, i2, seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, ?, D1, ?, D3> tensor, Index<D0> i0, Span ignored1, Index<D2> i2) {
-        return index(tensor, i0, seq(tensor.d1()), i2, seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<?, D1, ?, D3>> tensor, Index<D0> i0, Span ignored1, Index<D2> i2) {
+        return index(tensor, i0, seq(tensor.shape().d1()), i2, seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, ?, ?, ?, D3> tensor, Index<D0> i0, Index<D1> i1, Index<D2> i2) {
-        return index(tensor, i0, i1, i2, seq(tensor.d3()));
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<?, ?, ?, D3>> tensor, Index<D0> i0, Index<D1> i1, Index<D2> i2) {
+        return index(tensor, i0, i1, i2, seq(tensor.shape().d3()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, D0, D1, D2, ?> tensor, Span ignored0, Span ignored1, Span ignored2, Index<D3> i3) {
-        return index(tensor, seq(tensor.d0()), seq(tensor.d1()), seq(tensor.d2()), i3);
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ? extends Shape<D0, D1, D2, ?>> tensor, Span ignored0, Span ignored1, Span ignored2, Index<D3> i3) {
+        return index(tensor, seq(tensor.shape().d0()), seq(tensor.shape().d1()), seq(tensor.shape().d2()), i3);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, D0, D1, D2, D3> index(
-        Tensor<T, ?, ?, ?, ?> tensor, Index<D0> i0, Index<D1> i1, Index<D2> i2, Index<D3> i3) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> index(
+        Tensor<T, ?> tensor, Index<D0> i0, Index<D1> i1, Index<D2> i2, Index<D3> i3) {
         return operation("index")
                    .inputs(tensor)
                    .outputs(prototype(tensor.type(),
@@ -1470,34 +1465,34 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
 
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> List<Tensor<T, D0, N, U, U>> batch(
-        Tensor<T, D0, D1, U, U> tensor, int batchSize) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> List<Tensor<T, Shape<D0, N, U, U>>> batch(
+        Tensor<T, Shape<D0, D1, U, U>> tensor, int batchSize) {
         return batch(tensor, ArrayFire::n, batchSize);
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, BDT extends Num<BDT>> List<Tensor<T, D0, BDT, U, U>> batch(
-        Tensor<T, D0, D1, D2, D3> tensor, Function<Integer, BDT> type, int batchSize) {
-        var results = new ArrayList<Tensor<T, D0, BDT, U, U>>();
-        var d0Seq = seq(tensor.d0());
-        for (int i = 0; i < tensor.d1().size(); i += batchSize) {
-            var computedD1Size = Math.min(batchSize, tensor.d1().size() - i);
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, BDT extends Num<BDT>> List<Tensor<T, Shape<D0, BDT, U, U>>> batch(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, Function<Integer, BDT> type, int batchSize) {
+        var results = new ArrayList<Tensor<T, Shape<D0, BDT, U, U>>>();
+        var d0Seq = seq(tensor.shape().d0());
+        for (int i = 0; i < tensor.shape().d1().size(); i += batchSize) {
+            var computedD1Size = Math.min(batchSize, tensor.shape().d1().size() - i);
             var slice = index(tensor, d0Seq, seq(i, i + computedD1Size - 1));
-            results.add(slice.reshape(shape(tensor.d0(), type.apply(computedD1Size))));
+            results.add(slice.reshape(shape(tensor.shape().d0(), type.apply(computedD1Size))));
         }
         return results;
     }
 
     @SuppressWarnings({"unchecked", "rawtypes"})
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> tileAs(
-        Tensor<T, D0, D1, D2, D3> tensor, Shape<OD0, OD1, OD2, OD3> newShape) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>, NS extends Shape<?, ?, ?, ?>> Tensor<T, NS> tileAs(
+        Tensor<T, S> tensor, NS newShape) {
         if (newShape.capacity() % tensor.shape().capacity() != 0) {
             throw new IllegalArgumentException(
                 String.format("Can't tile perfectly from %s to %s", tensor.shape(), newShape));
         }
-        int d0ratio = newShape.d0().size() / tensor.d0().size();
-        int d1ratio = newShape.d1().size() / tensor.d1().size();
-        int d2ratio = newShape.d2().size() / tensor.d2().size();
-        int d3ratio = newShape.d3().size() / tensor.d3().size();
+        int d0ratio = newShape.d0().size() / tensor.shape().d0().size();
+        int d1ratio = newShape.d1().size() / tensor.shape().d1().size();
+        int d2ratio = newShape.d2().size() / tensor.shape().d2().size();
+        int d3ratio = newShape.d3().size() / tensor.shape().d3().size();
         return operation("tile")
                    .inputs(tensor)
                    .outputs(prototype(tensor.type(), newShape))
@@ -1507,39 +1502,38 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
     }
 
     @SuppressWarnings({"unchecked", "rawtypes"})
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<ST, OD0, OD1, OD2, OD3> sumAs(
-        Tensor<T, D0, D1, D2, D3> input, Shape<OD0, OD1, OD2, OD3> newShape) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, S extends Shape<?, ?, ?, ?>, NS extends Shape<?, ?, ?, ?>> Tensor<ST, NS> sumAs(
+        Tensor<T, S> input, NS newShape) {
         // I think there is a nicer way to do this in at most two operations.
         Tensor result = input;
-        if (newShape.d0() != input.d0()) {
+        if (newShape.d0() != input.shape().d0()) {
             if (newShape.d0().size() != 1)
                 throw new IllegalArgumentException("Can't sum over D0 from " + input.shape() + " to " + newShape);
             result = sum(result);
         }
-        if (newShape.d1() != input.d1()) {
+        if (newShape.d1() != input.shape().d1()) {
             if (newShape.d1().size() != 1)
                 throw new IllegalArgumentException("Can't sum over D1 from " + input.shape() + " to " + newShape);
             result = sum(result);
         }
-        if (newShape.d2() != input.d2()) {
+        if (newShape.d2() != input.shape().d2()) {
             if (newShape.d2().size() != 1)
                 throw new IllegalArgumentException("Can't sum over D2 from " + input.shape() + " to " + newShape);
             result = sum(result);
         }
-        if (newShape.d3() != input.d3()) {
+        if (newShape.d3() != input.shape().d3()) {
             if (newShape.d3().size() != 1)
                 throw new IllegalArgumentException("Can't sum over D3 from " + input.shape() + " to " + newShape);
             result = sum(result);
         }
-        return ((Tensor<ST, ?, ?, ?, ?>) result).reshape(newShape);
+        return reshape(((Tensor<ST, ?>) result), newShape);
     }
 
-    public static <T extends DataType<?, ?>> Tensor<T, N, U, U, U> flatten(Tensor<T, ?, ?, ?, ?> tensor) {
+    public static <T extends DataType<?, ?>> Tensor<T, Shape<N, U, U, U>> flatten(Tensor<T, ?> tensor) {
         return reshape(tensor, shape(tensor.shape().capacity()));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> flip(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Tensor<T, S> flip(Tensor<T, S> tensor) {
         return operation("flip")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1548,31 +1542,33 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, FD3 extends Num<?>> Tensor<T, N, N, FD3, D3> convolve2(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, ?, ?, D2, FD3> filters) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, FD0 extends Num<FD0>, FD1 extends Num<FD1>, FD3 extends Num<FD3>> Tensor<T, Shape<N, N, FD3, D3>> convolve2(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, Tensor<T, Shape<FD0, FD1, D2, FD3>> filters) {
         return convolve2(tensor, filters, shape(1, 1), shape(0, 0), shape(1, 1));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, FD3 extends Num<?>> Tensor<T, N, N, FD3, D3> convolve2(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, ?, ?, D2, FD3> filters, Shape<?, ?, ?, ?> stride) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, FD0 extends Num<FD0>, FD1 extends Num<FD1>, FD3 extends Num<FD3>> Tensor<T, Shape<N, N, FD3, D3>> convolve2(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, Tensor<T, Shape<FD0, FD1, D2, FD3>> filters,
+        Shape<?, ?, ?, ?> stride) {
         return convolve2(tensor, filters, stride, shape(0, 0), shape(1, 1));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, FD3 extends Num<?>> Tensor<T, N, N, FD3, D3> convolve2(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, ?, ?, D2, FD3> filters, Shape<?, ?, ?, ?> stride,
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, FD0 extends Num<FD0>, FD1 extends Num<FD1>, FD3 extends Num<FD3>> Tensor<T, Shape<N, N, FD3, D3>> convolve2(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, Tensor<T, Shape<FD0, FD1, D2, FD3>> filters, Shape<?, ?, ?, ?> stride,
         Shape<?, ?, ?, ?> padding) {
         return convolve2(tensor, filters, stride, padding, shape(1, 1));
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>, FD3 extends Num<?>> Tensor<T, N, N, FD3, D3> convolve2(
-        Tensor<T, D0, D1, D2, D3> tensor, Tensor<T, ?, ?, D2, FD3> filters, Shape<?, ?, ?, ?> stride,
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>, FD0 extends Num<FD0>, FD1 extends Num<FD1>, FD3 extends Num<FD3>> Tensor<T, Shape<N, N, FD3, D3>> convolve2(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor, Tensor<T, Shape<FD0, FD1, D2, FD3>> filters, Shape<?, ?, ?, ?> stride,
         Shape<?, ?, ?, ?> padding, Shape<?, ?, ?, ?> dilation) {
         // TODO: CoPilot wrote this, needs tests.
-        var computedShape = shape(
-            n((tensor.d0().size() + 2 * padding.d0().size() - (filters.d0().size() - 1) * dilation.d0().size() - 1) /
-                  stride.d0().size() + 1),
-            n((tensor.d1().size() + 2 * padding.d1().size() - (filters.d1().size() - 1) * dilation.d1().size() - 1) /
-                  stride.d1().size() + 1), filters.d3(), tensor.d3());
+        var computedShape = shape(n((tensor.shape().d0().size() + 2 * padding.d0().size() -
+                                         (filters.shape().d0().size() - 1) * dilation.d0().size() - 1) /
+                                        stride.d0().size() + 1),
+            n((tensor.shape().d1().size() + 2 * padding.d1().size() -
+                   (filters.shape().d1().size() - 1) * dilation.d1().size() - 1) / stride.d1().size() + 1),
+            filters.shape().d3(), tensor.shape().d3());
         return operation("convolve2")
                    .inputs(tensor, filters)
                    .outputs(prototype(tensor.type(), computedShape))
@@ -1588,8 +1584,8 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
     /**
      * L2 norm.
      */
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, U, D1, D2, D3> norm(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<U, D1, D2, D3>> norm(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         var mul = mul(tensor, tensor);
         var sum = sum(mul);
         return sqrt(sum);
@@ -1598,27 +1594,27 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
     /**
      * Normalize by dividing by the L2 norm.
      */
-    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<ST, D0, D1, D2, D3> normalize(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <ST extends DataType<?, ?>, T extends DataType<?, ST>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<ST, Shape<D0, D1, D2, D3>> normalize(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return div(cast(tensor, tensor.type().sumType()), norm(tensor).tileAs(tensor.shape()));
     }
 
     /**
      * Center by subtracting the average.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Tensor<T, D0, D1, D2, D3> center(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>> Tensor<T, Shape<D0, D1, D2, D3>> center(
+        Tensor<T, Shape<D0, D1, D2, D3>> tensor) {
         return sub(tensor, mean(tensor).tileAs(tensor));
     }
 
     // svd
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>> SvdResult<T, D0, D1> svd(
-        Tensor<T, D0, D1, U, U> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>> SvdResult<T, D0, D1> svd(
+        Tensor<T, Shape<D0, D1, U, U>> tensor) {
         var trio = operation("svd")
                        .inputs(tensor)
-                       .outputs(prototype(tensor.type(), shape(tensor.d0(), tensor.d0())),
-                           prototype(tensor.type(), shape(tensor.d0())),
-                           prototype(tensor.type(), shape(tensor.d1(), tensor.d1())))
+                       .outputs(prototype(tensor.type(), shape(tensor.shape().d0(), tensor.shape().d0())),
+                           prototype(tensor.type(), shape(tensor.shape().d0())),
+                           prototype(tensor.type(), shape(tensor.shape().d1(), tensor.shape().d1())))
                        .operation((u, s, v) -> arrayfire_h.af_svd(u, s, v, tensor.dereference()))
                        .build();
         return new SvdResult<>(trio.left(), trio.middle(), trio.right());
@@ -1627,34 +1623,34 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
     /**
      * Computes the covariance matrix of the given matrix.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>> Tensor<T, D0, D0, U, U> cov(
-        Tensor<T, D0, D1, U, U> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>> Tensor<T, Shape<D0, D0, U, U>> cov(
+        Tensor<T, Shape<D0, D1, U, U>> tensor) {
         return tidy(() -> {
             var subMean = sub(tensor, mean(tensor, D1).tileAs(tensor));
-            var matrix = matmul(subMean, subMean.transpose());
-            return div(matrix, constant(matrix.type(), matrix.shape(), tensor.d1().size() - 1.0f));
+            var matrix = matmul(subMean, transpose(subMean));
+            return div(matrix, constant(matrix.type(), matrix.shape(), tensor.shape().d1().size() - 1.0f));
         });
     }
 
     /**
      * Computes the ZCA whitening matrix of the given matrix.
      */
-    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<?>> Tensor<T, D0, D0, U, U> zca(
-        Tensor<T, D0, D1, U, U> tensor) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>> Tensor<T, Shape<D0, D0, U, U>> zca(
+        Tensor<T, Shape<D0, D1, U, U>> tensor) {
         return tidy(() -> {
             var cov = cov(tensor);
             var svd = svd(cov);
             var invSqrtS = diag(div(constant(svd.s().type(), svd.s().shape(), 1.0f),
                 sqrt(add(svd.s(), constant(svd.s().type(), svd.s().shape(), 1e-5f)))));
-            return matmul(svd.u(), matmul(invSqrtS, svd.u().transpose()));
+            return matmul(svd.u(), matmul(invSqrtS, transpose(svd.u())));
         });
     }
 
     /**
      * Inverts the given matrix.
      */
-    public static <T extends DataType<?, ?>, D extends Num<?>> Tensor<T, D, D, U, U> inverse(
-        Tensor<T, D, D, U, U> tensor) {
+    public static <T extends DataType<?, ?>, D extends Num<D>> Tensor<T, Shape<D, D, U, U>> inverse(
+        Tensor<T, Shape<D, D, U, U>> tensor) {
         return operation("inverse")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1663,8 +1659,8 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
     }
 
     // TODO: Add uncropped version.
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>> Tensor<T, D0, D1, U, U> rotate(
-        Tensor<T, D0, D1, U, U> tensor, float angle, InterpolationType interpolationType) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>> Tensor<T, Shape<D0, D1, U, U>> rotate(
+        Tensor<T, Shape<D0, D1, U, U>> tensor, float angle, InterpolationType interpolationType) {
         return operation("rotate")
                    .inputs(tensor)
                    .outputs(prototype(tensor))
@@ -1674,15 +1670,15 @@ public static <T extends DataType<?, T>, D0 extends Num<?>, D1 extends Num<?>, D
                    .build();
     }
 
-    public static <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, ND0 extends Num<?>, ND1 extends Num<?>> Tensor<T, ND0, ND1, U, U> scale(
-        Tensor<T, D0, D1, U, U> tensor, ND0 nd0, ND1 nd1, InterpolationType interpolationType) {
+    public static <T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>, ND0 extends Num<ND0>, ND1 extends Num<ND1>> Tensor<T, Shape<ND0, ND1, U, U>> scale(
+        Tensor<T, Shape<D0, D1, U, U>> tensor, ND0 nd0, ND1 nd1, InterpolationType interpolationType) {
         return operation("scale")
                    .inputs(tensor)
                    .outputs(prototype(tensor.type(), shape(nd0, nd1)))
-                   .operation(
-                       ptr -> arrayfire_h.af_scale(ptr, tensor.dereference(), (float) nd0.size() / tensor.d0().size(),
-                           (float) nd1.size() / tensor.d1().size(), nd0.size(), nd1.size(), interpolationType.code()))
-                   .grads((result, grads) -> scale(grads, tensor.d0(), tensor.d1(), interpolationType))
+                   .operation(ptr -> arrayfire_h.af_scale(ptr, tensor.dereference(),
+                       (float) nd0.size() / tensor.shape().d0().size(), (float) nd1.size() / tensor.shape().d1().size(),
+                       nd0.size(), nd1.size(), interpolationType.code()))
+                   .grads((result, grads) -> scale(grads, tensor.shape().d0(), tensor.shape().d1(), interpolationType))
                    .build();
     }
 
@@ -1805,12 +1801,12 @@ public static U u(int value) {
         return U;
     }
 
-    public static <T extends Tensor<?, ?, ?, ?, ?>> T grads(Tensor<?, ?, ?, ?, ?> loss, T tensor) {
+    public static <T extends Tensor<?, ?>> T grads(Tensor<?, ?> loss, T tensor) {
         var graph = new Graph(scope().operations());
         return graph.grads(loss, tensor);
     }
 
-    public static void optimize(Tensor<?, ?, ?, ?, ?> loss) {
+    public static void optimize(Tensor<?, ?> loss) {
         var graph = new Graph(scope().operations());
         graph.optimize(loss);
     }
diff --git a/arrayfire/ArrayFireTest.java b/arrayfire/ArrayFireTest.java
index 2202974..73db212 100644
--- a/arrayfire/ArrayFireTest.java
+++ b/arrayfire/ArrayFireTest.java
@@ -132,7 +132,7 @@ public void sortIndex() {
     public void permutationIndex() {
         af.tidy(() -> {
             var arr = af.create(1, 2, 3, 4, 5, 6, 7, 8).reshape(2, 4);
-            var permutation = af.permutation(arr.d1());
+            var permutation = af.permutation(arr.shape().d1());
             var shuffled = af.index(arr, af.span(), permutation);
             var data = af.data(shuffled);
             assertArrayEquals(new int[]{5, 6, 1, 2, 7, 8, 3, 4}, data.java());
@@ -143,7 +143,7 @@ public void permutationIndex() {
     public void transpose() {
         af.tidy(() -> {
             var arr = af.create(new float[]{1, 2, 3, 4}).reshape(2, 2);
-            var transpose = arr.transpose();
+            var transpose = af.transpose(arr);
             assertArrayEquals(new float[]{1, 3, 2, 4}, af.data(transpose).java(), 1E-5f);
         });
     }
@@ -172,7 +172,7 @@ public void matmul() {
         af.tidy(() -> {
             var left = af.create(new float[]{1, 2, 3, 4}).reshape(a(2), b(2));
             var right = af.create(new float[]{1, 2, 3, 4, 5, 6}).reshape(a(2), c(3));
-            var result = af.matmul(left.transpose(), right);
+            var result = af.matmul(af.transpose(left), right);
             assertArrayEquals(new float[]{5, 11, 11, 25, 17, 39}, data(result).java(), 1E-5f);
         });
     }
@@ -182,7 +182,7 @@ public void matmulS32() {
         af.tidy(() -> {
             var left = af.create(new float[]{1, 2, 3, 4}).reshape(a(2), b(2));
             var right = af.create(new float[]{1, 2, 3, 4, 5, 6}).reshape(a(2), c(3));
-            var result = af.matmul(left.transpose(), right);
+            var result = af.matmul(af.transpose(left), right);
             assertArrayEquals(new float[]{5, 11, 11, 25, 17, 39}, data(result).java(), 1E-5f);
         });
     }
@@ -292,7 +292,7 @@ public void mulScalar() {
     public void min() {
         af.tidy(() -> {
             var data = af.create(new float[]{-5, 12, 0, 1});
-            var result = data.min();
+            var result = af.min(data);
             assertArrayEquals(new float[]{-5}, af.data(result).java(), 1e-5f);
         });
     }
@@ -384,9 +384,9 @@ public void index4D() {
             var data = af
                            .create(new float[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
                            .reshape(2, 2, 2, 2);
-            Tensor<F32, A, B, C, D> result = af.index(data, af.seq(af.create(0).castshape(af::a)),
-                af.seq(af.create(1).castshape(af::b)), af.seq(af.create(0).castshape(af::c)),
-                af.seq(af.create(1).castshape(af::d)));
+            Tensor<F32, Shape<A, B, C, D>> result = af.index(data, af.seq(af.create(0).reshape(af.a(1))),
+                af.seq(af.create(1).reshape(af.b(1))), af.seq(af.create(0).reshape(af.c(1))),
+                af.seq(af.create(1).reshape(af.d(1))));
             assertArrayEquals(new float[]{11}, af.data(result).java(), 1E-5f);
         });
     }
@@ -554,7 +554,7 @@ public void graph() {
         af.tidy(() -> {
             var left = af.create(new float[]{1, 2, 3, 4}).reshape(a(2), b(2));
             var right = af.create(new float[]{1, 2, 3, 4, 5, 6}).reshape(a(2), c(3));
-            var leftT = left.transpose();
+            var leftT = af.transpose(left);
             var matmul = af.matmul(leftT, right);
             var softmax = af.softmax(matmul);
             var sum = af.sum(matmul);
diff --git a/arrayfire/GradFunction.java b/arrayfire/GradFunction.java
index b507754..b496a1e 100644
--- a/arrayfire/GradFunction.java
+++ b/arrayfire/GradFunction.java
@@ -1,8 +1,5 @@
 package arrayfire;
 
-import arrayfire.DataType;
-import arrayfire.Tensor;
-import arrayfire.TensorPair;
 import arrayfire.numbers.Num;
 
 import java.util.List;
@@ -10,15 +7,13 @@
 @FunctionalInterface
 interface GradFunction {
 
-    List<Tensor<?, ?, ?, ?, ?>> grads(Tensor<?, ?, ?, ?, ?> resultGrads);
+    List<Tensor<?, ?>> grads(Tensor<?, ?> resultGrads);
 
-    interface Unary<RT extends DataType<?, ?>, RD0 extends Num<?>, RD1 extends Num<?>, RD2 extends Num<?>, RD3 extends Num<?>, I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>> {
-        Tensor<I0T, I0D0, I0D1, I0D2, I0D3> grads(Tensor<RT, RD0, RD1, RD2, RD3> result,
-                                                  Tensor<RT, RD0, RD1, RD2, RD3> grads);
+    interface Unary<RT extends Tensor<?, ?>, IT extends Tensor<?, ?>> {
+        IT grads(RT result, RT grads);
     }
 
-    interface Binary<RT extends DataType<?, ?>, RD0 extends Num<?>, RD1 extends Num<?>, RD2 extends Num<?>, RD3 extends Num<?>, I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>, I1T extends DataType<?, ?>, I1D0 extends Num<?>, I1D1 extends Num<?>, I1D2 extends Num<?>, I1D3 extends Num<?>> {
-        TensorPair<I0T, I0D0, I0D1, I0D2, I0D3, I1T, I1D0, I1D1, I1D2, I1D3> grads(
-            Tensor<RT, RD0, RD1, RD2, RD3> result, Tensor<RT, RD0, RD1, RD2, RD3> grads);
+    interface Binary<RT extends Tensor<?, ?>, I0T extends Tensor<?, ?>, I1T extends Tensor<?, ?>> {
+        TensorPair<I0T, I1T> grads(RT result, RT grads);
     }
 }
diff --git a/arrayfire/Graph.java b/arrayfire/Graph.java
index 29a0a77..1ab3783 100644
--- a/arrayfire/Graph.java
+++ b/arrayfire/Graph.java
@@ -74,7 +74,7 @@ public void optimize(Tensor loss) {
         }
     }
 
-    public <T extends Tensor<?, ?, ?, ?, ?>> T grads(Tensor loss, T tensor) {
+    public <T extends Tensor<?, ?>> T grads(Tensor loss, T tensor) {
         var grads = grads(loss, new Tensor[]{tensor});
         return grads.get(tensor);
     }
@@ -153,7 +153,7 @@ void put(Tensor tensor, Tensor grads) {
         }
 
         @SuppressWarnings("unchecked")
-        public <T extends Tensor<?, ?, ?, ?, ?>> T get(T tensor) {
+        public <T extends Tensor<?, ?>> T get(T tensor) {
             return (T) gradsByTensor.get(tensor);
         }
     }
diff --git a/arrayfire/ImaxResult.java b/arrayfire/ImaxResult.java
index 7508c87..70a54cf 100644
--- a/arrayfire/ImaxResult.java
+++ b/arrayfire/ImaxResult.java
@@ -1,7 +1,5 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
-
-public record ImaxResult<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>>(
-    Tensor<T, D0, D1, D2, D3> values, Tensor<U32, D0, D1, D2, D3> indices) {
+public record ImaxResult<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>>(Tensor<T, S> values,
+                                                                                Tensor<U32, S> indices) {
 }
\ No newline at end of file
diff --git a/arrayfire/Index.java b/arrayfire/Index.java
index 7ab7ec9..24502ea 100644
--- a/arrayfire/Index.java
+++ b/arrayfire/Index.java
@@ -17,12 +17,12 @@ public class Index<D extends Num<D>> {
         ValueLayout.JAVA_BOOLEAN.withName("isSeq"), ValueLayout.JAVA_BOOLEAN.withName("isBatch"),
         MemoryLayout.paddingLayout(6));
 
-    private final Tensor<?, ?, ?, ?, ?> arr;
+    private final Tensor<?, ?> arr;
     private final Seq seq;
 
     private final Function<Integer, D> generator;
 
-    Index(Tensor<?, ?, ?, ?, ?> arr, Function<Integer, D> generator) {
+    Index(Tensor<?, ?> arr, Function<Integer, D> generator) {
         this.arr = arr;
         this.seq = null;
         this.generator = generator;
diff --git a/arrayfire/Operation.java b/arrayfire/Operation.java
index 7f42c4d..06ec72c 100644
--- a/arrayfire/Operation.java
+++ b/arrayfire/Operation.java
@@ -1,6 +1,5 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
 import arrayfire.utils.Functions;
 
 import java.lang.foreign.MemorySegment;
@@ -39,7 +38,7 @@ public void apply() {
         }
     }
 
-    public GradFunction grads() {
+    GradFunction grads() {
         return grads;
     }
 
@@ -56,14 +55,13 @@ public Nullary inputs() {
             return new Nullary();
         }
 
-        public <I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>> Unary<I0T, I0D0, I0D1, I0D2, I0D3> inputs(
-            Tensor<I0T, I0D0, I0D1, I0D2, I0D3> input) {
+        public <IT extends Tensor<?, ?>> Unary<IT> inputs(IT input) {
             operation.inputs.add(input);
             return new Unary<>();
         }
 
-        public <I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>, I1T extends DataType<?, ?>, I1D0 extends Num<?>, I1D1 extends Num<?>, I1D2 extends Num<?>, I1D3 extends Num<?>> Binary<I0T, I0D0, I0D1, I0D2, I0D3, I1T, I1D0, I1D1, I1D2, I1D3> inputs(
-            Tensor<I0T, I0D0, I0D1, I0D2, I0D3> left, Tensor<I1T, I1D0, I1D1, I1D2, I1D3> right) {
+        public <I0T extends DataType<?, ?>, I0S extends Shape<?, ?, ?, ?>, I1T extends DataType<?, ?>, I1S extends Shape<?, ?, ?, ?>> Binary<Tensor<I0T, I0S>, Tensor<I1T, I1S>> inputs(
+            Tensor<I0T, I0S> left, Tensor<I1T, I1S> right) {
             operation.inputs.add(left);
             operation.inputs.add(right);
             return new Binary<>();
@@ -71,49 +69,48 @@ public Nullary inputs() {
 
         public class Nullary {
 
-            public <OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Single<OT, OD0, OD1, OD2, OD3> outputs(
-                Prototype<OT, OD0, OD1, OD2, OD3> prototype) {
+            public <OT extends DataType<?, ?>, OS extends Shape<?, ?, ?, ?>> Single<Tensor<OT, OS>> outputs(
+                Prototype<OT, OS> prototype) {
                 operation.outputs.add(new Tensor<>(prototype));
                 return new Single<>();
             }
 
-            public class Single<OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> {
+            public class Single<OT extends Tensor<?, ?>> {
 
-                public Single<OT, OD0, OD1, OD2, OD3> operation(Function<MemorySegment, Integer> function) {
+                public Single<OT> operation(Function<MemorySegment, Integer> function) {
                     operation.apply = (outputs) -> af.handleStatus(() -> function.apply(outputs.getFirst().segment()));
                     return this;
                 }
 
                 @SuppressWarnings("unchecked")
-                public Tensor<OT, OD0, OD1, OD2, OD3> build() {
+                public OT build() {
                     af.scope().register(operation);
-                    return (Tensor<OT, OD0, OD1, OD2, OD3>) operation.outputs.getFirst();
+                    return (OT) operation.outputs.getFirst();
                 }
             }
         }
 
-        public class Unary<I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>> {
+        public class Unary<IT extends Tensor<?, ?>> {
 
-            public <OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> None outputs() {
+            public None outputs() {
                 return new None();
             }
 
-            public <OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Single<OT, OD0, OD1, OD2, OD3> outputs(
-                Prototype<OT, OD0, OD1, OD2, OD3> prototype) {
+            public <OT extends DataType<?, ?>, OS extends Shape<?, ?, ?, ?>> Single<Tensor<OT, OS>> outputs(
+                Prototype<OT, OS> prototype) {
                 operation.outputs.add(new Tensor<>(prototype));
                 return new Single<>();
             }
 
-            public <O0T extends DataType<?, ?>, O0D0 extends Num<?>, O0D1 extends Num<?>, O0D2 extends Num<?>, O0D3 extends Num<?>, O1T extends DataType<?, ?>, O1D0 extends Num<?>, O1D1 extends Num<?>, O1D2 extends Num<?>, O1D3 extends Num<?>> Pair<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3> outputs(
-                Prototype<O0T, O0D0, O0D1, O0D2, O0D3> left, Prototype<O1T, O1D0, O1D1, O1D2, O1D3> right) {
+            public <O0T extends DataType<?, ?>, O0S extends Shape<?, ?, ?, ?>, O1T extends DataType<?, ?>, O1S extends Shape<?, ?, ?, ?>> Pair<Tensor<O0T, O0S>, Tensor<O1T, O1S>> outputs(
+                Prototype<O0T, O0S> left, Prototype<O1T, O1S> right) {
                 operation.outputs.add(new Tensor<>(left));
                 operation.outputs.add(new Tensor<>(right));
                 return new Pair<>();
             }
 
-            public <O0T extends DataType<?, ?>, O0D0 extends Num<?>, O0D1 extends Num<?>, O0D2 extends Num<?>, O0D3 extends Num<?>, O1T extends DataType<?, ?>, O1D0 extends Num<?>, O1D1 extends Num<?>, O1D2 extends Num<?>, O1D3 extends Num<?>, O2T extends DataType<?, ?>, O2D0 extends Num<?>, O2D1 extends Num<?>, O2D2 extends Num<?>, O2D3 extends Num<?>> Trio<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3, O2T, O2D0, O2D1, O2D2, O2D3> outputs(
-                Prototype<O0T, O0D0, O0D1, O0D2, O0D3> left, Prototype<O1T, O1D0, O1D1, O1D2, O1D3> middle,
-                Prototype<O2T, O2D0, O2D1, O2D2, O2D3> right) {
+            public <O0T extends DataType<?, ?>, O0S extends Shape<?, ?, ?, ?>, O1T extends DataType<?, ?>, O1S extends Shape<?, ?, ?, ?>, O2T extends DataType<?, ?>, O2S extends Shape<?, ?, ?, ?>> Trio<Tensor<O0T, O0S>, Tensor<O1T, O1S>, Tensor<O2T, O2S>> outputs(
+                Prototype<O0T, O0S> left, Prototype<O1T, O1S> middle, Prototype<O2T, O2S> right) {
                 operation.outputs.add(new Tensor<>(left));
                 operation.outputs.add(new Tensor<>(middle));
                 operation.outputs.add(new Tensor<>(right));
@@ -133,51 +130,47 @@ public Operation build() {
                 }
             }
 
-            public class Single<OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> {
+            public class Single<OT extends Tensor<?, ?>> {
 
-                public Single<OT, OD0, OD1, OD2, OD3> operation(Function<MemorySegment, Integer> function) {
+                public Single<OT> operation(Function<MemorySegment, Integer> function) {
                     operation.apply = (outputs) -> af.handleStatus(() -> function.apply(outputs.getFirst().segment()));
                     return this;
                 }
 
                 @SuppressWarnings("unchecked")
-                public Single<OT, OD0, OD1, OD2, OD3> grads(
-                    GradFunction.Unary<OT, OD0, OD1, OD2, OD3, I0T, I0D0, I0D1, I0D2, I0D3> unaryGradFunction) {
+                public Single<OT> grads(GradFunction.Unary<OT, IT> unaryGradFunction) {
                     operation.grads = (grads) -> {
-                        var inputGrad = unaryGradFunction.grads(operation.outputs.getFirst(), (Tensor) grads);
+                        var inputGrad = unaryGradFunction.grads((OT) operation.outputs.getFirst(), (OT) grads);
                         return List.of(inputGrad);
                     };
                     return this;
                 }
 
                 @SuppressWarnings("unchecked")
-                public Tensor<OT, OD0, OD1, OD2, OD3> build() {
+                public OT build() {
                     af.scope().register(operation);
-                    return (Tensor<OT, OD0, OD1, OD2, OD3>) operation.outputs.getFirst();
+                    return (OT) operation.outputs.getFirst();
                 }
             }
 
-            public class Pair<O0T extends DataType<?, ?>, O0D0 extends Num<?>, O0D1 extends Num<?>, O0D2 extends Num<?>, O0D3 extends Num<?>, O1T extends DataType<?, ?>, O1D0 extends Num<?>, O1D1 extends Num<?>, O1D2 extends Num<?>, O1D3 extends Num<?>> {
+            public class Pair<O0T extends Tensor<?, ?>, O1T extends Tensor<?, ?>> {
 
-                public Pair<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3> operation(
-                    Functions.Function2<MemorySegment, MemorySegment, Integer> function) {
+                public Pair<O0T, O1T> operation(Functions.Function2<MemorySegment, MemorySegment, Integer> function) {
                     operation.apply = (outputs) -> af.handleStatus(
                         () -> function.apply(outputs.getFirst().segment(), outputs.get(1).segment()));
                     return this;
                 }
 
                 @SuppressWarnings("unchecked")
-                public TensorPair<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3> build() {
+                public TensorPair<O0T, O1T> build() {
                     af.scope().register(operation);
-                    return new TensorPair<>(
-                        (Tensor<O0T, O0D0, O0D1, O0D2, O0D3>) operation.outputs.getFirst(),
-                        (Tensor<O1T, O1D0, O1D1, O1D2, O1D3>) operation.outputs.get(1));
+                    return new TensorPair<>((O0T) operation.outputs.getFirst(), (O1T) operation.outputs.get(1));
                 }
             }
 
-            public class Trio<O0T extends DataType<?, ?>, O0D0 extends Num<?>, O0D1 extends Num<?>, O0D2 extends Num<?>, O0D3 extends Num<?>, O1T extends DataType<?, ?>, O1D0 extends Num<?>, O1D1 extends Num<?>, O1D2 extends Num<?>, O1D3 extends Num<?>, O2T extends DataType<?, ?>, O2D0 extends Num<?>, O2D1 extends Num<?>, O2D2 extends Num<?>, O2D3 extends Num<?>> {
+            public class Trio<O0T extends Tensor<?, ?>, O1T extends Tensor<?, ?>, O2T extends Tensor<?, ?>> {
 
-                public Trio<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3, O2T, O2D0, O2D1, O2D2, O2D3> operation(
+                public Trio<O0T, O1T, O2T> operation(
                     Functions.Function3<MemorySegment, MemorySegment, MemorySegment, Integer> function) {
                     operation.apply = (outputs) -> af.handleStatus(
                         () -> function.apply(outputs.getFirst().segment(), outputs.get(1).segment(),
@@ -186,45 +179,42 @@ public Trio<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3, O2T, O2D0,
                 }
 
                 @SuppressWarnings("unchecked")
-                public TensorTrio<O0T, O0D0, O0D1, O0D2, O0D3, O1T, O1D0, O1D1, O1D2, O1D3, O2T, O2D0, O2D1, O2D2, O2D3> build() {
+                public TensorTrio<O0T, O1T, O2T> build() {
                     af.scope().register(operation);
-                    return new TensorTrio<>(
-                        (Tensor<O0T, O0D0, O0D1, O0D2, O0D3>) operation.outputs.getFirst(),
-                        (Tensor<O1T, O1D0, O1D1, O1D2, O1D3>) operation.outputs.get(1),
-                        (Tensor<O2T, O2D0, O2D1, O2D2, O2D3>) operation.outputs.get(2));
+                    return new TensorTrio<>((O0T) operation.outputs.getFirst(), (O1T) operation.outputs.get(1),
+                        (O2T) operation.outputs.get(2));
                 }
             }
         }
 
-        public class Binary<I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>, I1T extends DataType<?, ?>, I1D0 extends Num<?>, I1D1 extends Num<?>, I1D2 extends Num<?>, I1D3 extends Num<?>> {
+        public class Binary<I0T extends Tensor<?, ?>, I1T extends Tensor<?, ?>> {
 
-            public <OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Single<OT, OD0, OD1, OD2, OD3> outputs(
-                Prototype<OT, OD0, OD1, OD2, OD3> prototype) {
+            public <OT extends DataType<?, ?>, OS extends Shape<?, ?, ?, ?>> Single<Tensor<OT, OS>> outputs(
+                Prototype<OT, OS> prototype) {
                 operation.outputs.add(new Tensor<>(prototype));
                 return new Single<>();
             }
 
-            public class Single<OT extends DataType<?, ?>, OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> {
+            public class Single<OT extends Tensor<?, ?>> {
 
-                public Single<OT, OD0, OD1, OD2, OD3> operation(Function<MemorySegment, Integer> function) {
+                public Single<OT> operation(Function<MemorySegment, Integer> function) {
                     operation.apply = (outputs) -> af.handleStatus(() -> function.apply(outputs.getFirst().segment()));
                     return this;
                 }
 
                 @SuppressWarnings("unchecked")
-                public Single<OT, OD0, OD1, OD2, OD3> grads(
-                    GradFunction.Binary<OT, OD0, OD1, OD2, OD3, I0T, I0D0, I0D1, I0D2, I0D3, I1T, I1D0, I1D1, I1D2, I1D3> binaryGradFunction) {
+                public Single<OT> grads(GradFunction.Binary<OT, I0T, I1T> binaryGradFunction) {
                     operation.grads = (grads) -> {
-                        var inputGrad = binaryGradFunction.grads(operation.outputs.getFirst(), (Tensor) grads);
+                        var inputGrad = binaryGradFunction.grads((OT) operation.outputs.getFirst(), (OT) grads);
                         return List.of(inputGrad.left(), inputGrad.right());
                     };
                     return this;
                 }
 
                 @SuppressWarnings("unchecked")
-                public Tensor<OT, OD0, OD1, OD2, OD3> build() {
+                public OT build() {
                     af.scope().register(operation);
-                    return (Tensor<OT, OD0, OD1, OD2, OD3>) operation.outputs.getFirst();
+                    return (OT) operation.outputs.getFirst();
                 }
             }
         }
diff --git a/arrayfire/Optimizer.java b/arrayfire/Optimizer.java
index 6d03c82..2ed7970 100644
--- a/arrayfire/Optimizer.java
+++ b/arrayfire/Optimizer.java
@@ -1,8 +1,6 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
+public interface Optimizer<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> {
 
-public interface Optimizer<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> {
-
-    public void optimize(Params<T, D0, D1, D2, D3> params, Tensor<T, D0, D1, D2, D3> gradients);
+    public void optimize(Params<T, S> params, Tensor<T, S> gradients);
 }
diff --git a/arrayfire/Params.java b/arrayfire/Params.java
index d51e12a..090b17d 100644
--- a/arrayfire/Params.java
+++ b/arrayfire/Params.java
@@ -6,16 +6,16 @@
 /**
  * A variable with an optimizer.
  */
-public class Params<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> extends Variable<T, D0, D1, D2, D3> {
+public class Params<T extends DataType<?, ?>, S extends Shape<?, ? ,? ,?>> extends Variable<T, S> {
 
-    private final Optimizer<T, D0, D1, D2, D3> optimizer;
+    private final Optimizer<T, S> optimizer;
 
-    public Params(T type, Shape<D0, D1, D2, D3> shape, OptimizerProvider optimizerProvider) {
+    public Params(T type, S shape, OptimizerProvider optimizerProvider) {
         super(type, shape);
         this.optimizer = optimizerProvider.get();
     }
 
-    public void optimize(Tensor<T, D0, D1, D2, D3> gradients) {
+    public void optimize(Tensor<T, S> gradients) {
         if (optimizer == null) {
             throw new IllegalStateException("Attempting to optimize params but no optimizer is provided.");
         }
diff --git a/arrayfire/Prototype.java b/arrayfire/Prototype.java
index 59fb6a1..84b2a75 100644
--- a/arrayfire/Prototype.java
+++ b/arrayfire/Prototype.java
@@ -1,7 +1,4 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
-
-public record Prototype<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>>(
-    T type, Shape<D0, D1, D2, D3> shape) {
+public record Prototype<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>>(T type, S shape) {
 }
diff --git a/arrayfire/Shape.java b/arrayfire/Shape.java
index 0365f7f..61f10c1 100644
--- a/arrayfire/Shape.java
+++ b/arrayfire/Shape.java
@@ -7,7 +7,7 @@
 import java.util.function.Function;
 
 
-public record Shape<D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>>(D0 d0, D1 d1, D2 d2, D3 d3) {
+public record Shape<D0 extends Num<D0>, D1 extends Num<D1>, D2 extends Num<D2>, D3 extends Num<D3>>(D0 d0, D1 d1, D2 d2, D3 d3) {
 
     public int capacity() {
         return d0.size() * d1.size() * d2.size() * d3.size();
diff --git a/arrayfire/SortIndexResult.java b/arrayfire/SortIndexResult.java
index eef02cd..4710a50 100644
--- a/arrayfire/SortIndexResult.java
+++ b/arrayfire/SortIndexResult.java
@@ -1,7 +1,5 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
-
-public record SortIndexResult<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>>(
-        Tensor<T, D0, D1, D2, D3> values, Tensor<U32, D0, D1, D2, D3> indices) {
+public record SortIndexResult<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>>(Tensor<T, S> values,
+                                                                                     Tensor<U32, S> indices) {
 }
diff --git a/arrayfire/SvdResult.java b/arrayfire/SvdResult.java
index cb6affe..645332d 100644
--- a/arrayfire/SvdResult.java
+++ b/arrayfire/SvdResult.java
@@ -3,8 +3,6 @@
 import arrayfire.numbers.Num;
 import arrayfire.numbers.U;
 
-public record SvdResult<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>>(
-        Tensor<T, D0, D0, U, U> u,
-        Tensor<T, D0, U, U, U> s,
-        Tensor<T, D1, D1, U, U> vt) {
+public record SvdResult<T extends DataType<?, ?>, D0 extends Num<D0>, D1 extends Num<D1>>(
+    Tensor<T, Shape<D0, D0, U, U>> u, Tensor<T, Shape<D0, U, U, U>> s, Tensor<T, Shape<D1, D1, U, U>> vt) {
 }
diff --git a/arrayfire/Tensor.java b/arrayfire/Tensor.java
index d698aba..d400bbd 100644
--- a/arrayfire/Tensor.java
+++ b/arrayfire/Tensor.java
@@ -8,21 +8,20 @@
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.ValueLayout;
-import java.util.function.Function;
 
-public class Tensor<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> implements MemoryContainer {
+public class Tensor<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> implements MemoryContainer {
 
     // Contains a single device pointer.
     public static final AddressLayout LAYOUT = ValueLayout.ADDRESS;
     private final T type;
-    private final Shape<D0, D1, D2, D3> shape;
+    private final S shape;
     private final MemorySegment segment;
 
-    public Tensor(Prototype<T, D0, D1, D2, D3> prototype) {
+    public Tensor(Prototype<T, S> prototype) {
         this(prototype.type(), prototype.shape());
     }
 
-    Tensor(T type, Shape<D0, D1, D2, D3> shape) {
+    Tensor(T type, S shape) {
         this.type = type;
         this.shape = shape;
         this.segment = Arena.ofAuto().allocate(LAYOUT);
@@ -39,31 +38,16 @@ public MemorySegment dereference() {
         return segment.get(LAYOUT, 0L);
     }
 
-    public D0 d0() {
-        return shape.d0();
-    }
-
-    public D1 d1() {
-        return shape.d1();
-    }
-
-    public D2 d2() {
-        return shape.d2();
-    }
-
-    public D3 d3() {
-        return shape.d3();
-    }
 
     public int capacity() {
         return shape.capacity();
     }
 
-    public Shape<D0, D1, D2, D3> shape() {
+    public S shape() {
         return shape;
     }
 
-    public Prototype<T, D0, D1, D2, D3> prototype() {
+    public Prototype<T, S> prototype() {
         return new Prototype<>(type, shape);
     }
 
@@ -76,65 +60,42 @@ public String toString() {
         return "AfTensor{" + "type=" + type + ", shape=" + shape + '}';
     }
 
-    public Tensor<T, D1, D0, D2, D3> transpose() {
-        return af.transpose(this);
-    }
-
-    public <OD0 extends Num<?>> Tensor<T, OD0, D1, D2, D3> castshape(Function<Integer, OD0> d0) {
-        return af.castshape(this, d0);
-    }
-
-    public <OD0 extends Num<?>, OD1 extends Num<?>> Tensor<T, OD0, OD1, D2, D3> castshape(Function<Integer, OD0> d0,
-                                                                                          Function<Integer, OD1> d1) {
-        return af.castshape(this, d0, d1);
-    }
-
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>> Tensor<T, OD0, OD1, OD2, D3> castshape(
-        Function<Integer, OD0> d0, Function<Integer, OD1> d1, Function<Integer, OD2> d2) {
-        return af.castshape(this, d0, d1, d2);
-    }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> castshape(
-        Function<Integer, OD0> d0, Function<Integer, OD1> d1, Function<Integer, OD2> d2, Function<Integer, OD3> d3) {
-        return af.castshape(this, d0, d1, d2, d3);
-    }
-
-    public Tensor<T, N, U, U, U> reshape(int d0) {
+    public Tensor<T, Shape<N, U, U, U>> reshape(int d0) {
         return af.reshape(this, af.shape(d0));
     }
 
-    public Tensor<T, N, N, U, U> reshape(int d0, int d1) {
+    public Tensor<T, Shape<N, N, U, U>> reshape(int d0, int d1) {
         return af.reshape(this, af.shape(d0, d1));
     }
 
-    public Tensor<T, N, N, N, U> reshape(int d0, int d1, int d2) {
+    public Tensor<T, Shape<N, N, N, U>> reshape(int d0, int d1, int d2) {
         return af.reshape(this, af.shape(d0, d1, d2));
     }
 
-    public Tensor<T, N, N, N, N> reshape(int d0, int d1, int d2, int d3) {
+    public Tensor<T, Shape<N, N, N, N>> reshape(int d0, int d1, int d2, int d3) {
         return af.reshape(this, af.shape(d0, d1, d2, d3));
     }
 
-    public <OD0 extends Num<?>> Tensor<T, OD0, U, U, U> reshape(OD0 d0) {
+    public <OD0 extends Num<OD0>> Tensor<T, Shape<OD0, U, U, U>> reshape(OD0 d0) {
         return af.reshape(this, af.shape(d0));
     }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>> Tensor<T, OD0, OD1, U, U> reshape(OD0 d0, OD1 d1) {
+    public <OD0 extends Num<OD0>, OD1 extends Num<OD1>> Tensor<T, Shape<OD0, OD1, U, U>> reshape(OD0 d0, OD1 d1) {
         return af.reshape(this, af.shape(d0, d1));
     }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>> Tensor<T, OD0, OD1, OD2, U> reshape(OD0 d0,
-                                                                                                            OD1 d1,
-                                                                                                            OD2 d2) {
+    public <OD0 extends Num<OD0>, OD1 extends Num<OD1>, OD2 extends Num<OD2>> Tensor<T, Shape<OD0, OD1, OD2, U>> reshape(
+        OD0 d0, OD1 d1, OD2 d2) {
         return af.reshape(this, af.shape(d0, d1, d2));
     }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> reshape(
+    public <OD0 extends Num<OD0>, OD1 extends Num<OD1>, OD2 extends Num<OD2>, OD3 extends Num<OD3>> Tensor<T, Shape<OD0, OD1, OD2, OD3>> reshape(
         OD0 d0, OD1 d1, OD2 d2, OD3 d3) {
         return af.reshape(this, af.shape(d0, d1, d2, d3));
     }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> reshape(
+    public <OD0 extends Num<OD0>, OD1 extends Num<OD1>, OD2 extends Num<OD2>, OD3 extends Num<OD3>> Tensor<T, Shape<OD0, OD1, OD2, OD3>> reshape(
         Shape<OD0, OD1, OD2, OD3> newShape) {
         return af.reshape(this, newShape);
     }
@@ -143,115 +104,77 @@ public void release() {
         af.release(this);
     }
 
-    Tensor<T, D0, D1, D2, D3> retain() {
+    Tensor<T, S> retain() {
         return af.retain(this);
     }
 
 
-    public Tensor<T, D0, D1, D2, D3> eval() {
+    public Tensor<T, S> eval() {
         return af.eval(this);
     }
 
-    public Tensor<T, U, D1, D2, D3> mean() {
-        return af.mean(this);
-    }
-
-    public Tensor<T, U, D1, D2, D3> mean(arrayfire.D0 dim) {
-        return af.mean(this, dim);
-    }
-
-    public Tensor<T, D0, U, D2, D3> mean(arrayfire.D1 dim) {
-        return af.mean(this, dim);
-    }
-
-    public Tensor<T, U, D1, D2, D3> median() {
-        return af.median(this);
-    }
-
-    public Tensor<T, U, D1, D2, D3> max() {
-        return af.max(this);
-    }
-
-    public Tensor<T, D0, U, D2, D3> max(arrayfire.D1 dim) {
-        return af.max(this, dim);
-    }
-
-    public Tensor<T, U, D1, D2, D3> min() {
-        return af.min(this);
-    }
-
-    public Tensor<T, D0, D1, D2, D3> clamp(Tensor<T, D0, D1, D2, D3> lo, Tensor<T, D0, D1, D2, D3> hi) {
+    public Tensor<T, S> clamp(Tensor<T, S> lo, Tensor<T, S> hi) {
         return af.clamp(this, lo, hi);
     }
 
-    public Tensor<T, D0, D1, D2, D3> relu() {
+    public Tensor<T, S> relu() {
         return af.relu(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> negate() {
+    public Tensor<T, S> negate() {
         return af.negate(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> exp() {
+    public Tensor<T, S> exp() {
         return af.exp(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> abs() {
+    public Tensor<T, S> abs() {
         return af.abs(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> sqrt() {
+    public Tensor<T, S> sqrt() {
         return af.sqrt(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> sigmoid() {
+    public Tensor<T, S> sigmoid() {
         return af.sigmoid(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> sparse(Storage storage) {
+    public Tensor<T, S> sparse(Storage storage) {
         return af.sparse(this, storage);
     }
 
-    public Tileable<T, D0, D1, D2, D3> tile() {
+    public Tileable<T, S> tile() {
         return new Tileable<>(this);
     }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> tileAs(
-        Tensor<T, OD0, OD1, OD2, OD3> newShapeTensor) {
+    public <NS extends Shape<?, ?, ?, ?>> Tensor<T, NS> tileAs(Tensor<T, NS> newShapeTensor) {
         return af.tileAs(this, newShapeTensor.shape());
     }
 
-    public <OD0 extends Num<?>, OD1 extends Num<?>, OD2 extends Num<?>, OD3 extends Num<?>> Tensor<T, OD0, OD1, OD2, OD3> tileAs(
-        Shape<OD0, OD1, OD2, OD3> newShape) {
+    public <NS extends Shape<?, ?, ?, ?>> Tensor<T, NS> tileAs(NS newShape) {
         return af.tileAs(this, newShape);
     }
 
-    public Tensor<T, N, U, U, U> flatten() {
+    public Tensor<T, Shape<N, U, U, U>> flatten() {
         return af.flatten(this);
     }
 
-    public Tensor<T, D0, D1, D2, D3> flip() {
+    public Tensor<T, S> flip() {
         return af.flip(this);
 
     }
 
-    public Tensor<T, D0, D1, D2, D3> move(Scope scope) {
+    public Tensor<T, S> move(Scope scope) {
         Scope.move(this, scope);
         return this;
     }
 
-    public <TN extends DataType<?, ?>> Tensor<TN, D0, D1, D2, D3> cast(TN t) {
+    public <TN extends DataType<?, ?>> Tensor<TN, S> cast(TN t) {
         return af.cast(this, t);
     }
 
-    /**
-     * Normalize by dividing by the L2 norm.
-     */
-
-    public Tensor<T, D0, D1, D2, D3> center() {
-        return af.center(this);
-    }
-
     @Override
     public void dispose() {
         release();
diff --git a/arrayfire/TensorPair.java b/arrayfire/TensorPair.java
index 2135819..2c2d81b 100644
--- a/arrayfire/TensorPair.java
+++ b/arrayfire/TensorPair.java
@@ -1,10 +1,5 @@
 package arrayfire;
 
-import arrayfire.DataType;
-import arrayfire.Tensor;
-import arrayfire.numbers.Num;
-
-public record TensorPair<I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>, I1T extends DataType<?, ?>, I1D0 extends Num<?>, I1D1 extends Num<?>, I1D2 extends Num<?>, I1D3 extends Num<?>>(
-    Tensor<I0T, I0D0, I0D1, I0D2, I0D3> left, Tensor<I1T, I1D0, I1D1, I1D2, I1D3> right) {
+public record TensorPair<L extends Tensor<?, ?>, R extends Tensor<?, ?>>(L left, R right) {
 
 }
diff --git a/arrayfire/TensorTrio.java b/arrayfire/TensorTrio.java
index 80fab9b..6104de0 100644
--- a/arrayfire/TensorTrio.java
+++ b/arrayfire/TensorTrio.java
@@ -1,11 +1,6 @@
 package arrayfire;
 
-import arrayfire.DataType;
-import arrayfire.Tensor;
-import arrayfire.numbers.Num;
-
-public record TensorTrio<I0T extends DataType<?, ?>, I0D0 extends Num<?>, I0D1 extends Num<?>, I0D2 extends Num<?>, I0D3 extends Num<?>, I1T extends DataType<?, ?>, I1D0 extends Num<?>, I1D1 extends Num<?>, I1D2 extends Num<?>, I1D3 extends Num<?>, I2T extends DataType<?, ?>, I2D0 extends Num<?>, I2D1 extends Num<?>, I2D2 extends Num<?>, I2D3 extends Num<?>>(
-    Tensor<I0T, I0D0, I0D1, I0D2, I0D3> left, Tensor<I1T, I1D0, I1D1, I1D2, I1D3> middle,
-    Tensor<I2T, I2D0, I2D1, I2D2, I2D3> right) {
+public record TensorTrio<T1 extends Tensor<?, ?>, T2 extends Tensor<?, ?>, T3 extends Tensor<?, ?>>(T1 left, T2 middle,
+                                                                                                    T3 right) {
 
 }
diff --git a/arrayfire/Tileable.java b/arrayfire/Tileable.java
index b350512..6e489fa 100644
--- a/arrayfire/Tileable.java
+++ b/arrayfire/Tileable.java
@@ -1,7 +1,4 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
-
-public record Tileable<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>>(
-        Tensor<T, D0, D1, D2, D3> tensor) {
+public record Tileable<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>>(Tensor<T, S> tensor) {
 }
diff --git a/arrayfire/TopKResult.java b/arrayfire/TopKResult.java
index 9926139..aa4b009 100644
--- a/arrayfire/TopKResult.java
+++ b/arrayfire/TopKResult.java
@@ -1,7 +1,5 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
-
-public record TopKResult<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>>(
-    Tensor<T, D0, D1, D2, D3> values, Tensor<U32, D0, D1, D2, D3> indices) {
+public record TopKResult<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>>(Tensor<T, S> values,
+                                                                                Tensor<U32, S> indices) {
 }
diff --git a/arrayfire/Variable.java b/arrayfire/Variable.java
index b4924be..92f85ae 100644
--- a/arrayfire/Variable.java
+++ b/arrayfire/Variable.java
@@ -1,17 +1,15 @@
 package arrayfire;
 
-import arrayfire.numbers.Num;
-
 /**
  * A variable with an optimizer.
  */
-public class Variable<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> extends Tensor<T, D0, D1, D2, D3> {
+public class Variable<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> extends Tensor<T, S> {
 
-    public Variable(T type, Shape<D0, D1, D2, D3> shape) {
+    public Variable(T type, S shape) {
         super(type, shape);
     }
 
-    public void set(Tensor<T, D0, D1, D2, D3> tensor) {
+    public void set(Tensor<T, S> tensor) {
         af.set(this, tensor);
     }
 }
diff --git a/arrayfire/optimizers/OptimizerProvider.java b/arrayfire/optimizers/OptimizerProvider.java
index 993d91b..7e6e00f 100644
--- a/arrayfire/optimizers/OptimizerProvider.java
+++ b/arrayfire/optimizers/OptimizerProvider.java
@@ -1,9 +1,9 @@
 package arrayfire.optimizers;
 
-import arrayfire.Optimizer;
 import arrayfire.DataType;
-import arrayfire.numbers.Num;
+import arrayfire.Optimizer;
+import arrayfire.Shape;
 
 public interface OptimizerProvider {
-    <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Optimizer<T, D0, D1, D2, D3> get();
+    <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Optimizer<T, S> get();
 }
diff --git a/arrayfire/optimizers/SGD.java b/arrayfire/optimizers/SGD.java
index 5c34635..128d7cd 100644
--- a/arrayfire/optimizers/SGD.java
+++ b/arrayfire/optimizers/SGD.java
@@ -1,10 +1,6 @@
 package arrayfire.optimizers;
 
-import arrayfire.Optimizer;
-import arrayfire.Params;
-import arrayfire.Tensor;
-import arrayfire.af;
-import arrayfire.DataType;
+import arrayfire.*;
 import arrayfire.numbers.Num;
 
 public class SGD implements OptimizerProvider {
@@ -20,14 +16,14 @@ public SGD learningRate(double learningRate) {
         return this;
     }
 
-    public <T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> Optimizer<T, D0, D1, D2, D3> get() {
+    public <T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> Optimizer<T, S> get() {
         return new SGDOptimizer<>();
     }
 
-    public class SGDOptimizer<T extends DataType<?, ?>, D0 extends Num<?>, D1 extends Num<?>, D2 extends Num<?>, D3 extends Num<?>> implements Optimizer<T, D0, D1, D2, D3> {
+    public class SGDOptimizer<T extends DataType<?, ?>, S extends Shape<?, ?, ?, ?>> implements Optimizer<T, S> {
 
         @Override
-        public void optimize(Params<T, D0, D1, D2, D3> params, Tensor<T, D0, D1, D2, D3> gradients) {
+        public void optimize(Params<T, S> params, Tensor<T, S> gradients) {
             params.set(af.sub(params, af.mul(gradients, learningRate)));
         }
     }