rm TensorCast

Author: Michael Abbott

Project.toml

@@ -1,7 +1,7 @@
 name = "SliceMap"
 uuid = "82cb661a-3f19-5665-9e27-df437c7e54c8"
 authors = ["Michael Abbott"]
-version = "0.1.2"
+version = "0.2"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -18,7 +18,6 @@ ForwardDiff = "0.10"
 JuliennedArrays = "0.2"
 MacroTools = "0.4, 0.5"
 StaticArrays = "0.10, 0.11, 0.12"
-TensorCast = "0.1, 0.2"
 Tracker = "0.2"
 Zygote = "0.4"
 ZygoteRules = "0.2"

src/SliceMap.jl

@@ -3,7 +3,7 @@ module SliceMap
 
 export mapcols, MapCols, maprows, slicemap, tmapcols, ThreadMapCols
 
-using MacroTools, TensorCast, JuliennedArrays
+using MacroTools, JuliennedArrays
 
 using Tracker
 using Tracker: TrackedMatrix, track, @grad, data
@@ -81,12 +81,11 @@ e.g. if `dims=(2,4)` then `f` must map matrices to matrices.
 The gradient is for Zygote only.
 """
 function slicemap(f::Function, A::AbstractArray{T,N}, args...; dims) where {T,N}
-    code = ntuple(d -> d in dims ? (:) : (*), N)
-    B = TensorCast.sliceview(A, code)
+    code = ntuple(d -> d in dims ? True() : False(), N)
+    B = JuliennedArrays.Slices(A, code...)
     C = [ f(slice, args...) for slice in B ]
-    TensorCast.glue(C, code)
+    JuliennedArrays.Align(C, code...)
 end
-# TODO switch to JuliennedArrays, then rm TensorCast dep
 
 #========== Forward, Static ==========#
 
@@ -153,23 +152,6 @@ end
 
 #========== Gradients for Zygote ==========#
 
-# @require Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-
-# @init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("zygote.jl")
-# Now using ZygoteRules instead, mapcols etc above.
-
-#= TensorCast =#
-# These could move there, TODO
-
-@adjoint TensorCast.sliceview(A::AbstractArray, code::Tuple) =
-    TensorCast.sliceview(A, code), Δ -> (TensorCast.glue(Δ, code), nothing)
-
-@adjoint TensorCast.red_glue(A::AbstractArray, code::Tuple) =
-    TensorCast.red_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)
-
-@adjoint TensorCast.copy_glue(A::AbstractArray, code::Tuple) =
-    TensorCast.copy_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)
-
 #= JuliennedArrays =#
 
 @adjoint JuliennedArrays.Slices(whole, along...) =

test/runtests.jl

@@ -1,7 +1,7 @@
 
 using SliceMap
 using Test
-using ForwardDiff, Tracker, Zygote, TensorCast, JuliennedArrays
+using ForwardDiff, Tracker, Zygote, JuliennedArrays
 
 Zygote.refresh()
 
@@ -35,10 +35,6 @@ Zygote.refresh()
     @test grad ≈ Zygote.gradient(m -> sum(sin, tmapcols(fun, m)), mat)[1]
     @test grad ≈ Zygote.gradient(m -> sum(sin, ThreadMapCols{3}(fun, m)), mat)[1]
 
-    tcm(mat) = @cast out[i,j] := fun(mat[:,j])[i]
-    @test res ≈ tcm(mat)
-    @test grad ≈ Zygote.gradient(m -> sum(sin, tcm(m)), mat)[1]
-
     jcols(f,m) = Align(map(f, Slices(m, True(), False())), True(), False())
     @test res ≈ jcols(fun, mat)
     @test grad ≈ Zygote.gradient(m -> sum(sin, jcols(fun, m)), mat)[1]
@@ -61,10 +57,6 @@ end
     @test grad ≈ Zygote.gradient(m -> sum(sin, mapcols(fun, m)), mat)[1]
     @test grad ≈ Zygote.gradient(m -> sum(sin, MapCols{3}(fun, m)), mat)[1]
 
-    # tcm3(mat) = @cast out[_,j] := fun(mat[:,j]) # changed here too
-    # @test res ≈ tcm3(mat)
-    # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm3(m)), mat)[1]
-
 end
 @testset "columns -> matrix" begin
 
@@ -83,10 +75,6 @@ end
     @test grad ≈ Zygote.gradient(m -> sum(sin, mapcols(fun, m)), mat)[1]
     @test grad ≈ Zygote.gradient(m -> sum(sin, MapCols{3}(fun, m)), mat)[1]
 
-    # tcm4(mat) = @cast out[i⊗i′,j] := fun(mat[:,j])[i,i′]  i:3
-    # @test res ≈ tcm4(mat)
-    # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm4(m)), mat)[1]
-
 end
 @testset "columns w args" begin
 
@@ -105,10 +93,6 @@ end
     @test grad ≈ Zygote.gradient(m -> sum(sin, mapcols(fun, m, 5)), mat)[1]
     @test grad ≈ Zygote.gradient(m -> sum(sin, MapCols{3}(fun, m, 5)), mat)[1]
 
-    # tcm5(mat) = @cast out[i,j] := fun(mat[:,j], 5)[i]
-    # @test res ≈ tcm5(mat)
-    # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm5(m)), mat)[1]
-
 end
 @testset "rows" begin
 
@@ -122,10 +106,6 @@ end
     @test grad ≈ Tracker.gradient(m -> sum(sin, maprows(fun, m)), mat)[1]
     @test grad ≈ Zygote.gradient(m -> sum(sin, maprows(fun, m)), mat)[1]
 
-    # tcm2(mat) = @cast out[i,j] := fun(mat[i,:])[j]
-    # @test res ≈ tcm2(mat)
-    # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm2(m)), mat)[1]
-
     jrows(f,m) = Align(map(f, Slices(m, False(), True())), False(), True())
     @test res ≈ jrows(fun, mat)
     @test grad ≈ Zygote.gradient(m -> sum(sin, jrows(fun, m)), mat)[1]