code: def get_neighbor(
      idx: jax.Array,
      mesh: jax.sharding.Mesh,
      axis_name: str,
      *,
      direction: str,
  ) -> tuple[jax.Array, ...]:
      """Helper function that computes the mesh indices of a neighbor."""
      axis_names = mesh.axis_names
      which_axis = axis_names.index(axis_name)
      mesh_index = [
          idx if i == which_axis else lax.axis_index(a)
          for i, a in enumerate(axis_names)
      ]
      axis_size = lax.psum(1, axis_name)
      if direction == "right":
          next_idx = lax.rem(idx + 1, axis_size)
      else:
          left = idx - 1
          next_idx = jnp.where(left < 0, left + axis_size, left)
      mesh_index[which_axis] = next_idx
      return tuple(mesh_index)
  
file_path: jax/jax/experimental/pallas/ops/tpu/all_gather.py
imports: from __future__ import annotations
  import functools
  from collections.abc import Sequence
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental import shard_map
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: get_neighbor(idx, mesh, axis_name) and returns: tuple