Update base for Update on "[ET-VK] Implement generic reduction shader + mean, sum, amax, amin"

## Context

Introduce a generic shader to compute reduction along a single dim, and `keepdim = True`. With the generic shader template, `mean`, `sum`, `amin`, and `amax` can be implemented.

Differential Revision: [D64840504](https://our.internmc.facebook.com/intern/diff/D64840504/)

[ghstack-poisoned]

Author: SS-JIA