Fix batch dimension adjustment in shader indexing utils (#5399)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/5399

## Context

In the GLSL shader utility functions, the functions which use the axis map to obtain texture position or tensor index had a small bug where the size of the tensor dimension used to concatenate with the batch dimension was incorrectly retrieved as `sizes.w`. The correct way to retrieve this size is `sizes[axis_map[3]]`.

This diff fixes this bug.

ghstack-source-id: 242840432
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D62772113

fbshipit-source-id: 1285b1509ff652703c7ecf16dbb2cc30d6d70769