keras/src/layers/core/lambda_layer.py/Lambda/_add_batch
class Lambda:    def compute_output_shape(self, input_shape):
        if self._output_shape is None:
            # Leverage backend shape inference
            try:
                inputs = tree.map_shape_structure(
                    lambda x: backend.KerasTensor(x, dtype=self.compute_dtype),
                    input_shape,
                )
                output_spec = backend.compute_output_spec(self.call, inputs)
                return tree.map_structure(lambda x: x.shape, output_spec)
            except:
                raise NotImplementedError(
                    "We could not automatically infer the shape of "
                    "the Lambda's output. Please specify the `output_shape` "
                    "argument for this Lambda layer."
                )

        if callable(self._output_shape):
            return self._output_shape(input_shape)

        # Output shapes are passed directly and don't include batch dimension.
        batch_size = tree.flatten(input_shape)[0]

        def _add_batch(shape):
            return (batch_size,) + shape

        return tree.map_shape_structure(_add_batch, self._output_shape)