Batching into shared memory is deprecated, but essential for performance

I was doing some profiling of my data pipeline and found that the `Batch` transformation was a severe bottleneck. Here are the critical lines in [operations.py](https://github.com/google/grain/blob/main/grain/_src/python/operations.py):

```python
def stacking_function(*args):
first_arg = np.asanyarray(args[0])
shape, dtype = (len(args),) + first_arg.shape, first_arg.dtype
if not self._use_shared_memory or dtype.hasobject:
return np.stack(args)
return np.stack(args, out=SharedMemoryArray(shape, dtype=dtype)).metadata
```

I found that `self._use_shared_memory == True` iff you used the deprecated `grain.BatchOperation`, rather than the "recommended" `grain.Batch`. And what do you know, switching to `grain.BatchOperation` gave me a 3x increase in throughput! This matches up with my intuition, because in the `self._use_shared_memory == True` branch, there is only one copy that goes directly into shared memory. But in the `self._use_shared_memory == False` branch, the `np.stack` will induce one copy into private memory, and then the later [CopyNumPyArrayToSharedMemory transform](https://github.com/google/grain/blob/f463b5625558cc30308cccbe2102f2d1228f1af5/grain/_src/python/data_loader.py#L140) performs an explicit second copy into shared memory. It's not too surprising that adding another copy of all of the pipeline's data could slow things down significantly.

Here comes the real problem -- I want to use grain through airio, which doesn't go through the standard `DataLoader`, but the much more complex `lazy_dataset` API. In `lazy_dataset`, batching is done through a [different code path](https://github.com/google/grain/blob/f463b5625558cc30308cccbe2102f2d1228f1af5/grain/_src/python/lazy_dataset/transformations/batch.py#L31) that does not have an option to enable this optimization. It always batches into private memory, and then the [MultiprocessPrefetchLazyIterDataset](https://github.com/google/grain/blob/f463b5625558cc30308cccbe2102f2d1228f1af5/grain/_src/python/lazy_dataset/lazy_dataset.py#L513) does a second copy into shared memory.

I manually added a (slightly hacky) solution that enables batching directly into shared memory iff the batch operation is a parent of a `MultiprocessPrefetchLazyIterDataset`. Indeed, I saw a significant performance increase when using grain through airio. Is this something that could possibly be upstreamed into grain?