BmnGemStripMLPModel.cxx

Go to the documentation of this file.

#include "BmnGemStripMLPModel.h"
 
#include "FairLogger.h"
#include "TString.h"
#include "TSystem.h"
 
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <limits>
#include <map>
#include <mutex>
#include <set>
#include <sstream>
 
// Model file format (.bmnmlp, little-endian) ----------------------------------
// Written by Converter_KerasToBmnMLP.py, the version must be the same in both.
//
//   char[8]  magic "BMNMLP\0\0"
//   u32      format version (= kFormatVersion)
//   u32      architecture type (1 - residual MLP)
//   u32      input_dim, width, n_blocks, output_dim
//   u32      output activation (0 - linear, 1 - sigmoid)
//   f32      LayerNormalization epsilon
//   f32      LeakyReLU negative slope
//   str      arch_hash, source (.keras file name), created (date)       str = u32 length + bytes
//   u32      n_features, then n_features x str (feature names, the input order)
//   u32      n_params, then n_params x f32:
//              W_in[input_dim][width], b_in[width],
//              n_blocks x { ln_gamma[width], ln_beta[width], W1[width][width], b1[width], W2[width][width], b2[width]
//              }, ln_gamma[width], ln_beta[width], W_out[width][output_dim], b_out[output_dim]
//   u32      n_tests
//   f32      test tolerance (max absolute difference)
//   f32      n_tests x input_dim  (test inputs)
//   f32      n_tests x output_dim (Keras outputs)
//   u32      CRC32 of all preceding bytes
//
// Dense kernels are stored as in Keras: [in][out] (row of the input i = weights to all outputs).
//------------------------------------------------------------------------------
 
namespace
{
const char kMagic[8] = {'B', 'M', 'N', 'M', 'L', 'P', '\0', '\0'};
const uint32_t kFormatVersion = 1;
const uint32_t kArchResidualMLP = 1;
const char* const kBuiltinDefaultPath = "$VMCWORKDIR/parameters/gem/ANNModels/GEM_DefaultModel.bmnmlp";
 
// CRC32 (the same as zlib.crc32 in Python) ------------------------------------
uint32_t Crc32(const unsigned char* data, size_t n)
{
    static uint32_t table[256];
    static std::once_flag table_flag;
    std::call_once(table_flag, [] {
        for (uint32_t i = 0; i < 256; ++i) {
            uint32_t c = i;
            for (int k = 0; k < 8; ++k)
                c = (c & 1u) ? (0xEDB88320u ^ (c >> 1)) : (c >> 1);
            table[i] = c;
        }
    });
 
    uint32_t crc = 0xFFFFFFFFu;
    for (size_t i = 0; i < n; ++i)
        crc = table[(crc ^ data[i]) & 0xFFu] ^ (crc >> 8);
    return crc ^ 0xFFFFFFFFu;
}
 
bool IsLittleEndian()
{
    const uint32_t one = 1;
    unsigned char first;
    std::memcpy(&first, &one, 1);
    return first == 1;
}
 
// Reading the buffer with bounds checks ---------------------------------------
class Reader
{
  public:
    Reader(const unsigned char* data, size_t size)
        : fData(data)
        , fSize(size)
        , fPos(0)
    {}
 
    bool Bytes(void* dst, size_t n)
    {
        if (n > fSize - fPos)
            return false;
        std::memcpy(dst, fData + fPos, n);
        fPos += n;
        return true;
    }
    bool U32(uint32_t& v) { return Bytes(&v, sizeof(v)); }
    bool F32(float& v) { return Bytes(&v, sizeof(v)); }
    bool Floats(std::vector<float>& v, size_t n)
    {
        if (n > (fSize - fPos) / sizeof(float))
            return false;
        v.resize(n);
        return n == 0 || Bytes(v.data(), n * sizeof(float));
    }
    bool Str(std::string& s)
    {
        uint32_t len = 0;
        if (!U32(len) || len > (1u << 20) || len > fSize - fPos)
            return false;
        s.assign(reinterpret_cast<const char*>(fData + fPos), len);
        fPos += len;
        return true;
    }
    size_t Pos() const { return fPos; }
 
  private:
    const unsigned char* fData;
    size_t fSize;
    size_t fPos;
};
 
// Kernels ---------------------------------------------------------------------
// y[j] += sum_k xs[k] * W[row_k][j],  W[n_in][n_out] (Keras layout), row_k = rows[k] (kIndexed) or k.
// Explicit 4-float vectors (GCC/Clang vector extensions, SSE2/NEON baseline): the same speed at -O2 and -O3,
// without -march flags. 8 accumulators = 32 outputs in registers, rows of W are read sequentially.
// For each output the sum goes in the order of k, so PredictSplit() gives exactly the same result as Predict().
#if defined(__GNUC__)
typedef float v4f __attribute__((vector_size(16)));
 
inline v4f Load4(const float* p)
{
    v4f v;
    std::memcpy(&v, p, sizeof(v));
    return v;
}
 
inline void Store4(float* p, v4f v)
{
    std::memcpy(p, &v, sizeof(v));
}
#endif
 
template<bool kIndexed>
inline void AccumulateRows(float* __restrict y,
                           int n_out,
                           const float* __restrict W,
                           const float* __restrict xs,
                           const int* __restrict rows,
                           int n_rows)
{
    int j = 0;
#if defined(__GNUC__)
    for (; j + 32 <= n_out; j += 32) {
        v4f a0 = Load4(y + j), a1 = Load4(y + j + 4), a2 = Load4(y + j + 8), a3 = Load4(y + j + 12);
        v4f a4 = Load4(y + j + 16), a5 = Load4(y + j + 20), a6 = Load4(y + j + 24), a7 = Load4(y + j + 28);
        for (int k = 0; k < n_rows; ++k) {
            const float* w = W + static_cast<size_t>(kIndexed ? rows[k] : k) * n_out + j;
            const float s = xs[k];
            const v4f b = {s, s, s, s};
            a0 += b * Load4(w);
            a1 += b * Load4(w + 4);
            a2 += b * Load4(w + 8);
            a3 += b * Load4(w + 12);
            a4 += b * Load4(w + 16);
            a5 += b * Load4(w + 20);
            a6 += b * Load4(w + 24);
            a7 += b * Load4(w + 28);
        }
        Store4(y + j, a0);
        Store4(y + j + 4, a1);
        Store4(y + j + 8, a2);
        Store4(y + j + 12, a3);
        Store4(y + j + 16, a4);
        Store4(y + j + 20, a5);
        Store4(y + j + 24, a6);
        Store4(y + j + 28, a7);
    }
    for (; j + 4 <= n_out; j += 4) {
        v4f a0 = Load4(y + j);
        for (int k = 0; k < n_rows; ++k) {
            const float s = xs[k];
            const v4f b = {s, s, s, s};
            a0 += b * Load4(W + static_cast<size_t>(kIndexed ? rows[k] : k) * n_out + j);
        }
        Store4(y + j, a0);
    }
#endif
    for (; j < n_out; ++j) {
        float s = y[j];
        for (int k = 0; k < n_rows; ++k)
            s += xs[k] * W[static_cast<size_t>(kIndexed ? rows[k] : k) * n_out + j];
        y[j] = s;
    }
}
 
// Non-zero inputs (the segment features are often 0): x[0..n) -> rows = offset + i, vals = x[i]; returns the count
inline int CompactNonZero(const float* x, int n, int offset, int* rows, float* vals)
{
    int nnz = 0;
    for (int i = 0; i < n; ++i) {
        if (x[i] != 0.0f) {
            rows[nnz] = offset + i;
            vals[nnz] = x[i];
            ++nnz;
        }
    }
    return nnz;
}
 
// y = b + x * W,  W[n_in][n_out]
inline void Dense(const float* __restrict x,
                  int n_in,
                  const float* __restrict W,
                  const float* __restrict b,
                  int n_out,
                  float* __restrict y)
{
    std::memcpy(y, b, sizeof(float) * n_out);
    AccumulateRows<false>(y, n_out, W, x, nullptr, n_in);
}
 
inline void LeakyReLU(float* x, int n, float slope)
{
    for (int i = 0; i < n; ++i)
        x[i] = (x[i] > 0.0f) ? x[i] : slope * x[i];
}
 
// y = (x - mean) / sqrt(var + eps) * gamma + beta  (over the whole vector, biased variance as in Keras)
inline void LayerNorm(const float* __restrict x,
                      int n,
                      const float* __restrict gamma,
                      const float* __restrict beta,
                      float eps,
                      float* __restrict y)
{
    double sum = 0.0;
    for (int i = 0; i < n; ++i)
        sum += x[i];
    const double mean = sum / n;
    double sum2 = 0.0;
    for (int i = 0; i < n; ++i) {
        const double d = x[i] - mean;
        sum2 += d * d;
    }
    const double var = sum2 / n;
    const float inv = static_cast<float>(1.0 / std::sqrt(var + eps));
    const float fmean = static_cast<float>(mean);
    for (int i = 0; i < n; ++i)
        y[i] = (x[i] - fmean) * inv * gamma[i] + beta[i];
}
 
// Path with environment variables ($VMCWORKDIR, ~, ...) -> real path (ROOT: gSystem->ExpandPathName).
// ok = false if it cannot be expanded (an environment variable is not set); the path is returned as given then.
std::string ExpandPathName(const std::string& path, bool* ok = nullptr)
{
    TString expanded(path.c_str());
    const bool failed = gSystem->ExpandPathName(expanded);   // kTRUE in case of error
    if (ok)
        *ok = !failed;
    return failed ? path : std::string(expanded.Data());
}
 
// Shared models ---------------------------------------------------------------
struct SharedState
{
    std::mutex mutex;
    std::string default_path = kBuiltinDefaultPath;
    std::map<std::string, std::shared_ptr<const BmnGemStripMLPModel>> models;   // key - expanded path
    std::set<std::string> failed;                                               // paths that failed to load
    std::map<std::string, std::string> expanded;                                // path as given -> expanded path
};
 
SharedState& Shared()
{
    static SharedState state;
    return state;
}
//==============================================================================
}   // namespace
//------------------------------------------------------------------------------
 
BmnGemStripMLPModel::BmnGemStripMLPModel()
{
    Clear();
}
 
BmnGemStripMLPModel::BmnGemStripMLPModel(const std::string& path, bool verbose)
{
    Clear();
    Load(path, verbose);
}
 
BmnGemStripMLPModel::~BmnGemStripMLPModel() {}
 
void BmnGemStripMLPModel::Clear()
{
    fLoaded = false;
 
    fInputDim = 0;
    fWidth = 0;
    fNBlocks = 0;
    fOutputDim = 0;
    fOutputActivation = kLinearOutput;
    fLayerNormEps = 0.0f;
    fLeakyReLUSlope = 0.0f;
 
    fFilePath.clear();
    fArchHash.clear();
    fSourceName.clear();
    fCreated.clear();
    fFeatureNames.clear();
 
    std::vector<float>().swap(fParams);
    fOffDenseInW = fOffDenseInB = fOffBlocks = fBlockSize = 0;
    fOffFinalLnG = fOffFinalLnB = fOffDenseOutW = fOffDenseOutB = 0;
 
    fNTests = 0;
    fTestTolerance = 0.0f;
    std::vector<float>().swap(fTestInputs);
    std::vector<float>().swap(fTestOutputs);
}
//------------------------------------------------------------------------------
 
bool BmnGemStripMLPModel::Load(const std::string& path, bool verbose)
{
    Clear();
 
    bool expanded_ok = true;
    const std::string full_path = ExpandPathName(path, &expanded_ok);
 
    auto fail = [&](const std::string& reason) {
        if (verbose)
            LOG(error) << "BmnGemStripMLPModel: cannot load model " << full_path << ": " << reason;
        Clear();
        return false;
    };
 
    if (!IsLittleEndian())
        return fail("only little-endian platforms are supported");
 
    std::ifstream file(full_path, std::ios::binary | std::ios::ate);
    if (!file) {
        std::string reason = "file not found or not readable";
        if (!expanded_ok)
            reason += " (an environment variable in the path is not set)";
        return fail(reason);
    }
 
    const std::streamsize size = file.tellg();
    if (size <= 0)
        return fail("empty file");
    std::vector<unsigned char> buffer(static_cast<size_t>(size));
    file.seekg(0);
    if (!file.read(reinterpret_cast<char*>(buffer.data()), size))
        return fail("read error");
 
    std::string error;
    if (!Parse(buffer, error))
        return fail(error);
 
    fFilePath = full_path;
    fLoaded = true;
 
    // Self-test: the same results as Keras on the test samples from the file
    float max_diff = 0.0f;
    if (!RunSelfTest(&max_diff, false)) {
        std::ostringstream reason;
        reason << "self-test failed: max |output - Keras| = " << max_diff << " > tolerance " << fTestTolerance;
        return fail(reason.str());
    }
 
    return true;
}
//------------------------------------------------------------------------------
 
bool BmnGemStripMLPModel::Parse(const std::vector<unsigned char>& buffer, std::string& error)
{
    // Checksum: the last 4 bytes - CRC32 of everything before them
    if (buffer.size() < sizeof(kMagic) + 4 * sizeof(uint32_t)) {
        error = "file is too short";
        return false;
    }
    uint32_t crc_stored = 0;
    std::memcpy(&crc_stored, buffer.data() + buffer.size() - sizeof(uint32_t), sizeof(uint32_t));
    const size_t data_size = buffer.size() - sizeof(uint32_t);
    if (Crc32(buffer.data(), data_size) != crc_stored) {
        error = "checksum mismatch (the file is damaged or truncated)";
        return false;
    }
 
    Reader rd(buffer.data(), data_size);
    auto bad = [&](const std::string& what) {
        error = "format error: " + what;
        return false;
    };
 
    char magic[sizeof(kMagic)];
    if (!rd.Bytes(magic, sizeof(magic)) || std::memcmp(magic, kMagic, sizeof(kMagic)) != 0)
        return bad("not a .bmnmlp file (wrong signature)");
 
    uint32_t version = 0, arch = 0;
    if (!rd.U32(version))
        return bad("no version");
    if (version != kFormatVersion) {
        error = "unsupported format version " + std::to_string(version)
                + " (supported: " + std::to_string(kFormatVersion) + ")";
        return false;
    }
    if (!rd.U32(arch))
        return bad("no architecture type");
    if (arch != kArchResidualMLP) {
        error = "unsupported architecture type " + std::to_string(arch);
        return false;
    }
 
    uint32_t input_dim = 0, width = 0, n_blocks = 0, output_dim = 0, out_act = 0;
    float eps = 0.0f, slope = 0.0f;
    if (!rd.U32(input_dim) || !rd.U32(width) || !rd.U32(n_blocks) || !rd.U32(output_dim) || !rd.U32(out_act)
        || !rd.F32(eps) || !rd.F32(slope))
        return bad("truncated header");
 
    if (input_dim < 1 || input_dim > static_cast<uint32_t>(kMaxDim))
        return bad("input_dim = " + std::to_string(input_dim) + " is out of range");
    if (output_dim < 1 || output_dim > static_cast<uint32_t>(kMaxDim))
        return bad("output_dim = " + std::to_string(output_dim) + " is out of range");
    if (width < 1 || width > static_cast<uint32_t>(kMaxWidth))
        return bad("width = " + std::to_string(width) + " is out of range (max " + std::to_string(kMaxWidth) + ")");
    if (n_blocks > 10000)
        return bad("n_blocks = " + std::to_string(n_blocks) + " is out of range");
    if (out_act != kLinearOutput && out_act != kSigmoidOutput)
        return bad("unknown output activation " + std::to_string(out_act));
    if (!(eps > 0.0f) || !std::isfinite(eps) || !std::isfinite(slope))
        return bad("bad LayerNorm epsilon or LeakyReLU slope");
 
    std::string arch_hash, source, created;
    if (!rd.Str(arch_hash) || !rd.Str(source) || !rd.Str(created))
        return bad("truncated strings");
 
    uint32_t n_features = 0;
    if (!rd.U32(n_features))
        return bad("no feature count");
    if (n_features != input_dim)
        return bad("number of feature names (" + std::to_string(n_features) + ") != input_dim ("
                   + std::to_string(input_dim) + ")");
    std::vector<std::string> names(n_features);
    for (auto& name : names)
        if (!rd.Str(name))
            return bad("truncated feature names");
 
    // Weights
    const size_t w = width, in = input_dim, out = output_dim;
    const size_t block_size = 2 * w + 2 * (w * w + w);
    const size_t n_expected = in * w + w + n_blocks * block_size + 2 * w + w * out + out;
 
    uint32_t n_params = 0;
    if (!rd.U32(n_params))
        return bad("no parameter count");
    if (n_params != n_expected)
        return bad("number of parameters " + std::to_string(n_params) + ", expected " + std::to_string(n_expected));
    std::vector<float> params;
    if (!rd.Floats(params, n_params))
        return bad("truncated weights");
    for (float v : params)
        if (!std::isfinite(v))
            return bad("NaN or Inf in the weights");
 
    // Tests
    uint32_t n_tests = 0;
    float tolerance = 0.0f;
    if (!rd.U32(n_tests) || !rd.F32(tolerance))
        return bad("truncated test header");
    if (n_tests > 100000)
        return bad("too many test samples");
    std::vector<float> test_in, test_out;
    if (!rd.Floats(test_in, static_cast<size_t>(n_tests) * in)
        || !rd.Floats(test_out, static_cast<size_t>(n_tests) * out))
        return bad("truncated test samples");
 
    if (rd.Pos() != data_size)
        return bad("unexpected data at the end of the file");
 
    // Everything is read: fill the model
    fInputDim = static_cast<int>(input_dim);
    fWidth = static_cast<int>(width);
    fNBlocks = static_cast<int>(n_blocks);
    fOutputDim = static_cast<int>(output_dim);
    fOutputActivation = static_cast<int>(out_act);
    fLayerNormEps = eps;
    fLeakyReLUSlope = slope;
    fArchHash = arch_hash;
    fSourceName = source;
    fCreated = created;
    fFeatureNames = std::move(names);
    fParams = std::move(params);
 
    fOffDenseInW = 0;
    fOffDenseInB = fOffDenseInW + in * w;
    fOffBlocks = fOffDenseInB + w;
    fBlockSize = block_size;
    fOffFinalLnG = fOffBlocks + n_blocks * block_size;
    fOffFinalLnB = fOffFinalLnG + w;
    fOffDenseOutW = fOffFinalLnB + w;
    fOffDenseOutB = fOffDenseOutW + w * out;
 
    fNTests = static_cast<int>(n_tests);
    fTestTolerance = tolerance;
    fTestInputs = std::move(test_in);
    fTestOutputs = std::move(test_out);
    return true;
}
//------------------------------------------------------------------------------
 
void BmnGemStripMLPModel::ForwardFromFirstDense(float* h, float* output) const
{
    const int w = fWidth;
    const float* p = fParams.data();
    float t[kMaxWidth];
    float u[kMaxWidth];
 
    LeakyReLU(h, w, fLeakyReLUSlope);
 
    // Residual blocks: h = h + Dense2(LeakyReLU(Dense1(LayerNorm(h))))
    for (int iblock = 0; iblock < fNBlocks; ++iblock) {
        const float* blk = p + fOffBlocks + static_cast<size_t>(iblock) * fBlockSize;
        const float* ln_g = blk;
        const float* ln_b = ln_g + w;
        const float* W1 = ln_b + w;
        const float* b1 = W1 + static_cast<size_t>(w) * w;
        const float* W2 = b1 + w;
        const float* b2 = W2 + static_cast<size_t>(w) * w;
 
        LayerNorm(h, w, ln_g, ln_b, fLayerNormEps, t);
        Dense(t, w, W1, b1, w, u);
        LeakyReLU(u, w, fLeakyReLUSlope);
        Dense(u, w, W2, b2, w, t);
        for (int i = 0; i < w; ++i)
            h[i] += t[i];
    }
 
    // Final LayerNorm and Dense
    LayerNorm(h, w, p + fOffFinalLnG, p + fOffFinalLnB, fLayerNormEps, t);
    Dense(t, w, p + fOffDenseOutW, p + fOffDenseOutB, fOutputDim, output);
 
    if (fOutputActivation == kSigmoidOutput) {
        for (int j = 0; j < fOutputDim; ++j)
            output[j] = 1.0f / (1.0f + std::exp(-output[j]));
    }
}
//------------------------------------------------------------------------------
 
bool BmnGemStripMLPModel::Predict(const float* input, float* output) const
{
    if (!fLoaded || !input || !output)
        return false;
 
    const int w = fWidth;
    float h[kMaxWidth];
    int rows[kMaxDim];
    float vals[kMaxDim];
 
    // First Dense (zero inputs are skipped: the segment features are often 0)
    std::memcpy(h, fParams.data() + fOffDenseInB, sizeof(float) * w);
    const int nnz = CompactNonZero(input, fInputDim, 0, rows, vals);
    AccumulateRows<true>(h, w, fParams.data() + fOffDenseInW, vals, rows, nnz);
 
    ForwardFromFirstDense(h, output);
    return true;
}
 
std::vector<float> BmnGemStripMLPModel::Predict(const std::vector<float>& input) const
{
    if (!fLoaded || static_cast<int>(input.size()) != fInputDim)
        return {};
    std::vector<float> output(fOutputDim);
    if (!Predict(input.data(), output.data()))
        return {};
    return output;
}
//------------------------------------------------------------------------------
 
bool BmnGemStripMLPModel::PredictSplit(const float* common,
                                       int n_common,
                                       const float* specific,
                                       int n_rows,
                                       float* output) const
{
    if (!fLoaded || !output || n_rows < 0 || n_common < 0 || n_common > fInputDim)
        return false;
    const int n_spec = fInputDim - n_common;
    if ((n_common > 0 && !common) || (n_spec > 0 && n_rows > 0 && !specific))
        return false;
 
    const int w = fWidth;
    const float* W = fParams.data() + fOffDenseInW;
    int rows[kMaxDim];
    float vals[kMaxDim];
 
    // First Dense, common part (once for all rows)
    float base[kMaxWidth];
    std::memcpy(base, fParams.data() + fOffDenseInB, sizeof(float) * w);
    const int nnz_common = CompactNonZero(common, n_common, 0, rows, vals);
    AccumulateRows<true>(base, w, W, vals, rows, nnz_common);
 
    // Specific part and the rest of the network for each row
    float h[kMaxWidth];
    for (int r = 0; r < n_rows; ++r) {
        std::memcpy(h, base, sizeof(float) * w);
        const int nnz = CompactNonZero(specific + static_cast<size_t>(r) * n_spec, n_spec, n_common, rows, vals);
        AccumulateRows<true>(h, w, W, vals, rows, nnz);
        ForwardFromFirstDense(h, output + static_cast<size_t>(r) * fOutputDim);
    }
    return true;
}
//------------------------------------------------------------------------------
 
const float* BmnGemStripMLPModel::GetTestInput(int itest) const
{
    if (itest < 0 || itest >= fNTests)
        return nullptr;
    return fTestInputs.data() + static_cast<size_t>(itest) * fInputDim;
}
 
const float* BmnGemStripMLPModel::GetTestOutput(int itest) const
{
    if (itest < 0 || itest >= fNTests)
        return nullptr;
    return fTestOutputs.data() + static_cast<size_t>(itest) * fOutputDim;
}
 
bool BmnGemStripMLPModel::RunSelfTest(float* max_abs_diff, bool verbose) const
{
    if (max_abs_diff)
        *max_abs_diff = 0.0f;
    if (!fLoaded)
        return false;
 
    std::vector<float> out(fOutputDim);
    float max_diff = 0.0f;
    for (int itest = 0; itest < fNTests; ++itest) {
        Predict(GetTestInput(itest), out.data());
        const float* ref = GetTestOutput(itest);
        for (int j = 0; j < fOutputDim; ++j) {
            const float d = std::fabs(out[j] - ref[j]);
            if (!(d <= max_diff))   // NaN is also caught
                max_diff = std::isfinite(d) ? d : std::numeric_limits<float>::infinity();
        }
    }
    if (max_abs_diff)
        *max_abs_diff = max_diff;
 
    const bool ok = max_diff <= fTestTolerance;
    if (verbose)
        LOG(info) << "BmnGemStripMLPModel: self-test on " << fNTests << " samples: max |output - Keras| = " << max_diff
                  << " (tolerance " << fTestTolerance << ") -> " << (ok ? "OK" : "FAILED");
    return ok;
}
//------------------------------------------------------------------------------
 
int BmnGemStripMLPModel::GetFeatureIndex(const std::string& name) const
{
    auto it = std::find(fFeatureNames.begin(), fFeatureNames.end(), name);
    return (it == fFeatureNames.end()) ? -1 : static_cast<int>(it - fFeatureNames.begin());
}
 
bool BmnGemStripMLPModel::CheckFeatureNames(const std::vector<std::string>& expected, std::string* why) const
{
    if (fFeatureNames.size() != expected.size()) {
        if (why)
            *why = "model has " + std::to_string(fFeatureNames.size()) + " input features, expected "
                   + std::to_string(expected.size());
        return false;
    }
    for (size_t i = 0; i < expected.size(); ++i) {
        if (fFeatureNames[i] != expected[i]) {
            if (why)
                *why = "input feature " + std::to_string(i) + " is '" + fFeatureNames[i] + "', expected '" + expected[i]
                       + "'";
            return false;
        }
    }
    return true;
}
 
void BmnGemStripMLPModel::Print(std::ostream& os) const
{
    if (!fLoaded) {
        os << "BmnGemStripMLPModel: not loaded\n";
        return;
    }
    os << "BmnGemStripMLPModel: " << fFilePath << "\n"
       << "  residual MLP: input " << fInputDim << " -> width " << fWidth << " x " << fNBlocks << " blocks -> output "
       << fOutputDim << (fOutputActivation == kSigmoidOutput ? " (sigmoid)" : " (linear)") << "\n"
       << "  parameters: " << fParams.size() << ", LayerNorm eps = " << fLayerNormEps
       << ", LeakyReLU slope = " << fLeakyReLUSlope << "\n"
       << "  source: " << fSourceName << ", arch_hash: " << fArchHash << ", created: " << fCreated << "\n"
       << "  test samples: " << fNTests << " (tolerance " << fTestTolerance << ")\n"
       << "  features:";
    for (size_t i = 0; i < fFeatureNames.size(); ++i)
        os << (i % 16 == 0 ? "\n    " : " ") << fFeatureNames[i];
    os << "\n";
}
//------------------------------------------------------------------------------
 
std::shared_ptr<const BmnGemStripMLPModel> BmnGemStripMLPModel::GetShared(const std::string& path)
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
 
    // Cache key - the expanded path (the same file given by different strings is loaded once);
    // the expansion is done once for each path string (GetShared() is called for each point)
    const std::string& given = path.empty() ? state.default_path : path;
    auto iexp = state.expanded.find(given);
    if (iexp == state.expanded.end())
        iexp = state.expanded.emplace(given, ExpandPathName(given)).first;
    const std::string key = iexp->second;
 
    auto it = state.models.find(key);
    if (it != state.models.end())
        return it->second;
    if (state.failed.count(key))
        return nullptr;
 
    auto model = std::make_shared<BmnGemStripMLPModel>();
    if (!model->Load(key, true)) {
        LOG(warn) << "BmnGemStripMLPModel: the model " << key
                  << " is not available (this warning is printed once for the path)";
        state.failed.insert(key);
        return nullptr;
    }
    std::shared_ptr<const BmnGemStripMLPModel> shared = model;
    state.models[key] = shared;
    return shared;
}
 
void BmnGemStripMLPModel::SetDefaultPath(const std::string& path)
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
    state.default_path = path.empty() ? std::string(kBuiltinDefaultPath) : path;
}
 
std::string BmnGemStripMLPModel::GetDefaultPath()
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
    return state.default_path;
}
 
std::string BmnGemStripMLPModel::GetBuiltinDefaultPath()
{
    return kBuiltinDefaultPath;
}
 
void BmnGemStripMLPModel::ClearSharedCache()
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
    state.models.clear();   // models in use stay alive while someone holds them
    state.failed.clear();
    state.expanded.clear();
}