JuliaMLTools

using Flux
using Random
Random.seed!(12345)

function showdiff(f, n)
    a = rand(Float32, n, n) # shape (n,n)
    b = rand(Float32, n, n) # shape (n,n)
    c = cat(a, b; dims=3)   # shape (n,n,2)
    a_out = f(a)            # shape is (n,n)
    c_out = f(c)[:,:,1]     # shape is (n,n)