GatherV2 op missing support for batch_dims parameter

[pyu10055]

Please make sure that this is a bug. As per our
GitHub Policy,
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow.js installed from (npm or script link): 2.6.0
• TensorFlow.js version (use command below):
• Browser version:
• Tensorflow.js Converter Version: 2.6.0

Describe the current behavior
The GatherV2 op does not support batch_dims param as described in the python API
tf.raw_ops.GatherV2(
params, indices, axis, batch_dims=0, name=None
)

indices must be an integer tensor of any dimension (usually 0-D or 1-D). Produces an output tensor with shape params.shape[:axis] + indices.shape[batch_dims:] + params.shape[axis + 1:]

Ref model that is failing due to this missing support #4128
Describe the expected behavior

The GatherV2 API should support batch_dims param.

Standalone code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/CodePen/any notebook.

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

[pyu10055]

I have investigated the unstack op, the num parameter has to be same as the dimension of the axis, which means it is not needed.
The issue of #4128 is caused by GatherV2 op not supporting batch_dims parameter, I have verified with this fix, the model can be executed without errors.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No