keras/src/ops/numpy.py/Trace/compute_output_spec
class Trace:    def compute_output_spec(self, x):
        x_shape = list(x.shape)
        x_shape[self.axis1] = -1
        x_shape[self.axis2] = -1
        output_shape = list(filter((-1).__ne__, x_shape))
        output_dtype = backend.standardize_dtype(x.dtype)
        if output_dtype not in ("int64", "uint32", "uint64"):
            output_dtype = dtypes.result_type(output_dtype, "int32")
        return KerasTensor(output_shape, dtype=output_dtype)