keras/src/ops/core.py/Map/compute_output_spec
class Map:    def compute_output_spec(self, f, xs):
        x = xs[0]
        n = xs.shape[0]
        y = backend.compute_output_spec(f, x)

        def append_batch_axis(x):
            return KerasTensor(
                shape=(n,) + x.shape, dtype=x.dtype, sparse=x.sparse
            )

        y = tree.map_structure(append_batch_axis, y)
        return y