     }
 }
 
+public extension Tensor where R.Shape == Shape2D {
+    fileprivate func getElementShape() -> Shape1D {
+        return arrayToShape(Array(dynamicShape.dimensions.dropFirst()))
+    }
+}
+
 /// - TODO: Add conditional expressibility conformance in Swift 4
 
 public extension Tensor where R.DataType : Equatable {

            case is DataType.Type:
                return units[index] as! Element
            case is Tensor1D<DataType>.Type:
                let newShape: Shape1D = arrayToShape(Array(dynamicShape.dimensions.dropFirst()))
                let newUnits = ContiguousArray(units[getElementTensorRange(index: index)])
                return Tensor1D<DataType>(shape: newShape, units: newUnits) as! Element
            case is Tensor2D<DataType>.Type:
                let newShape: Shape2D = arrayToShape(Array(dynamicShape.dimensions.dropFirst()))
                let newUnits = ContiguousArray(units[getElementTensorRange(index: index)])
                return Tensor2D<DataType>(shape: newShape, units: newUnits) as! Element
            case is Tensor3D<DataType>.Type:
                let newShape: Shape3D = arrayToShape(Array(dynamicShape.dimensions.dropFirst()))
                let newUnits = ContiguousArray(units[getElementTensorRange(index: index)])
                return Tensor3D<DataType>(shape: newShape, units: newUnits) as! Element
            default:
                fatalError("Invalid element tensor type")
            }

    fileprivate func getElementTensorRange(index: Int) -> CountableRange<Int> {
        let elementTensorShape = dynamicShape.dropFirst()
        let contiguousIndex = unitIndex(fromIndex: index)
        return contiguousIndex..<contiguousIndex+elementTensorShape.contiguousSize
    }

    /// Get element tensor units at index
    fileprivate func getElementTensorUnits(index: Int) -> ContiguousArray<DataType> {
        let range = getElementTensorRange(index: index)
        return ContiguousArray(units[range])
    }

    /// Access the scalar element or element tensor at index
