keras/src/ops/numpy.py/Diff/compute_output_spec
class Diff:    def compute_output_spec(self, a):
        shape = list(a.shape)
        size = shape[self.axis]
        if size is not None:
            shape[self.axis] = builtins.max(size - self.n, 0)
        return KerasTensor(shape, dtype=a.dtype)