    def _infer_shapes(self, x: Tensor, indices: Tensor) -> List[IntVar]:
        """Infers shapes for batch_gather."""

        rank = len(indices._attrs["shape"])

        # TODO: remove this when we're sure we support non-static batch_gather
        x_shape_values = [var._attrs["values"][0] for var in x._attrs["shape"]]
        indices_shape = [var._attrs["values"][0] for var in indices._attrs["shape"]]
        for r in range(1, rank - 1):
            assert x_shape_values[r] == indices_shape[r]

        out_shapes = x._attrs["shape"][:]
        if rank <= 1:
            # Special case: gather happens along batch dimension
            out_shapes[0] = indices.shape()[0]
        out_shapes[rank - 1] = indices._attrs["shape"][-1]

        return out_shapes
Infers shapes for batch_gather.
python/aitemplate/compiler/ops/tensor/batch_gather.py