def _gather(params, indices, validate_indices=None, axis=None, batch_dims=0, name=None):
    'gather.'
    indices = ops.convert_to_tensor(indices, dtype_hint=np.int32)
    if (validate_indices is not None):
        raise NotImplementedError('Argument `validate_indices != None` is currently unimplemented.')
    if (batch_dims < 0):
        raise NotImplementedError('Negative `batch_dims` is currently unsupported.')
    if (axis is None):
        axis = batch_dims
    if (axis < 0):
        axis = (axis + len(params.shape))
    if JAX_MODE:
        take = (lambda params, indices: np.take(params, indices, axis=(axis - batch_dims)))
        take = functools.reduce((lambda g, f: f(g)), ([jax.vmap] * int(batch_dims)), take)
        return take(params, indices)
    params = ops.convert_to_tensor(params)
    res = np.array([np.take(params[i], indices[i], axis=(axis - batch_dims)) for i in np.ndindex(*params.shape[:batch_dims])])
    return np.reshape(res, ((params.shape[:axis] + indices.shape[batch_dims:]) + params.shape[(axis + 1):]))
gather.
def _gather(params, indices, validate_indices=None, axis=None, batch_dims=0, name=None):
    
    indices = ops.convert_to_tensor(indices, dtype_hint=np.int32)
    if (validate_indices is not None):
        raise NotImplementedError('Argument `validate_indices != None` is currently unimplemented.')
    if (batch_dims < 0):
        raise NotImplementedError('Negative `batch_dims` is currently unsupported.')
    if (axis is None):
        axis = batch_dims
    if (axis < 0):
        axis = (axis + len(params.shape))
    if JAX_MODE:
        take = (lambda params, indices: np.take(params, indices, axis=(axis - batch_dims)))
        take = functools.reduce((lambda g, f: f(g)), ([jax.vmap] * int(batch_dims)), take)
        return take(params, indices)
    params = ops.convert_to_tensor(params)
    res = np.array([np.take(params[i], indices[i], axis=(axis - batch_dims)) for i in np.ndindex(*params.shape[:batch_dims])])
    return np.reshape(res, ((params.shape[:axis] + indices.shape[batch_dims:]) + params.shape[(axis + 1):]))