Pass image as argument to distortionEffect metal shader

I spent a majority of the evening trying to figure out how to pass an image into an MSL shader. Here’s what I got working:

#define TAU 6.28318530717958647692528676656

[[stitchable]] float2 myDistortion(
    float2 position,
    float2 resolution,
    texture2d<half> source [[ texture(0)]]
    ) {
    float2 uv = position / resolution;
    //^^ Normalize coordinates

    constexpr sampler texSampler (mag_filter::linear, min_filter::linear);
    //^^ Create texture sampler
    half4 color = source.sample(texSampler, uv);
    //^^ sample the texture at the uv coordinates
    //... transform coordinates by color
    float angle = avg_color(color) * TAU;
    position = (uv + float2( cos(angle), sin(angle))) * resolution;
    return position;
}

Then, in the visualEffect modifier, give the distortion effect the resolution and the image. The image is converted into a texture2d automatically.

struct MyView: View {
    var image: Image?
    var enableDistortion: Bool

    var body: some View {
        image?
            .resizable()
            .aspectRatio(contentMode: .fit)
            .visualEffect { content, proxy in
                content
                    .distortionEffect(
                        ShaderLibrary.myDistortion(.float2(proxy.size), .image(image!)),
                        maxSampleOffset: .zero,
                        isEnabled: enableDistortion
                        )
            }
    }

}