GL_INVALID_OPERATION setting image2D in computeShader

I am using OpenGL ES 3.2 and the NVIDIA driver 361.00 on a Pixel C tablet with Tegra X1 GPU. I would like to use a compute shader to write data to a colour map and then later I will use some graphics shaders to display the image.

I already have this concept working using desktop GL and now I want to port to mobile. I am implementing the GL in Java rather than in native code. I extend GLSurfaceView and the GLSurfaceView.Renderer and then during the OnSurfaceCreated callback I initialise the shader programs and textures etc.

The compute shader compiles just fine without any errors:

#version 320 es
layout(binding = 0, rgba32f) uniform highp image2D colourMap;
layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;

void main()
{
    imageStore(colourMap, ivec2(gl_GlobalInvocationID.xy), vec4(1.0f, 0.0f, 0.0f, 1.0f));
}

And I initialise a texture

// Generate a 2D texture
GLES31.glGenTextures(1, colourMap, 0);
GLES31.glBindTexture(GLES31.GL_TEXTURE_2D, colourMap[0]);

// Set interpolation to nearest
GLES31.glTexParameteri(GLES31.GL_TEXTURE_2D, GLES31.GL_TEXTURE_MAG_FILTER, GLES31.GL_LINEAR);
GLES31.glTexParameteri(GLES31.GL_TEXTURE_2D, GLES31.GL_TEXTURE_MIN_FILTER, GLES31.GL_LINEAR);

// Create some dummy texture to begin with so we can see if it changes
float texData[] = new float[texWidth * texHeight * 4];
for (int j = 0; j < texHeight; j++)
{
    for (int i = 0; i < texWidth; i++)
    {
        // Set a few pixels in here...
    }
}
Buffer texDataBuffer = FloatBuffer.wrap(texData);
GLES31.glTexImage2D(GLES31.GL_TEXTURE_2D, 0, GLES31.GL_RGBA32F, texWidth, texHeight, 0, GLES31.GL_RGBA, GLES31.GL_FLOAT, texDataBuffer);

After this I set the image unit in the shader here:

GLES31.glUseProgram(idComputeShaderProgram);
int loc = GLES31.glGetUniformLocation(idComputeShaderProgram, "colourMap");
if (loc == -1) Log.e("Error", "Cannot locate variable");
GLES31.glUniform1i(loc, 0);

After every call to GL I check for errors using GLES31.glGetError() -- left out here for clarity. This final line is the only one which errors. The error code translates to GL_INVALID_OPERATION. As the documentation describes, glUniform1i can return this error under quite a few circumstances. I have checked all these and also believe that none of them apply. The shader compiles correctly and the program object is valid. The location of the variable is also valid. I have even used glGetActiveUniform() to get the type of the variable and it returns a type of 36941 which translates to GL_IMAGE_2D which I believe is an integer.

I don't know what I can be doing wrong. Am I missing an extension or something? Documentation for GLES in Java for Android is dire.