Add __getitem__ to backwards compatible shims.

PiperOrigin-RevId: 846196686

Author: Google-ML-Automation

jax/_src/interpreters/batching.py

@@ -660,6 +660,10 @@ def __setitem__(self, prim, batcher):
     def wrapped(axis_data, vals, dims, **params):
       return batcher(axis_data.size, axis_data.name, None, vals, dims, **params)
     fancy_primitive_batchers[prim] = wrapped
+
+  def __getitem__(self, prim):
+    return fancy_primitive_batchers[prim]
+
 axis_primitive_batchers = AxisPrimitiveBatchersProxy()
 
 # backwards compat shim. TODO: delete
@@ -675,6 +679,10 @@ def wrapped(axis_data, vals, dims, **params):
 
   def __delitem__(self, prim):
     del fancy_primitive_batchers[prim]
+
+  def __getitem__(self, prim):
+    return fancy_primitive_batchers[prim]
+
 primitive_batchers = PrimitiveBatchersProxy()