public:
  MOCK_METHOD(llvm::StringRef, name, (), (const, override));
  MOCK_METHOD(MlirOptimizationPassState, GetPassState,
              (const DeviceSet* device_set, const ConfigProto& config_proto,
               const Graph& graph,
               const FunctionLibraryDefinition& function_library),
              (const, override));
  MOCK_METHOD(Status, Run,
              (const std::string& function_name,

  MOCK_METHOD(absl::StatusOr<std::string>, SerializeExecutable,
              (const xla::LocalExecutable& executable), (override));
  MOCK_METHOD(absl::StatusOr<std::string>, BuildSerializedExecutable,
              (const XlaCompiler::Options& options,
               const XlaCompiler::CompilationResult& result),
              (override));
  MOCK_METHOD(absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>,

};

class MockDeviceCompilationProfiler : public DeviceCompilationProfiler {
 public:
  MOCK_METHOD(bool, ShouldCompileCluster,
              (const NameAttrList& function, DeviceCompileMode compile_mode,
               int64_t current_request_count),
              (override));
  MOCK_METHOD(Status, RegisterCompilation,
              (const NameAttrList& function, int64_t compile_time_us,

        backward_tensor = AddTensor(sizes[forward_tensor]);
        AddUse(op, backward_tensor);
      }
    }
    // We will also instrument the actual optimizations performed.
    MOCK_METHOD(void, ApplyRemat, (const RematSpec&));
  };
};

TEST_F(GreedyRematTest, MlpBasic) {
  StrictMock<MlpRemat> remat(std::vector<int>({1, 1, 1}));
  // (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum