diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
new file mode 100644
index 0000000..ed08691
--- /dev/null
+++ b/.devcontainer/Dockerfile
@@ -0,0 +1,46 @@
+FROM hexpm/elixir:1.20.2-erlang-29.0.2-debian-trixie-20260623-slim
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV LANG=C.UTF-8
+ENV LC_ALL=C.UTF-8
+ENV CC=gcc-14
+ENV CXX=g++-14
+
+RUN apt-get update && \
+  apt-get dist-upgrade --yes && \
+  apt-get install --yes --no-install-recommends \
+    build-essential \
+    ca-certificates \
+    cmake \
+    curl \
+    g++-14 \
+    gcc-14 \
+    git \
+    inotify-tools \
+    ninja-build \
+    openssh-client \
+    pkg-config \
+    python3 \
+    python3-venv \
+    sudo \
+    zsh && \
+  update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 140 && \
+  update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 140 && \
+  update-alternatives --install /usr/bin/cc cc /usr/bin/gcc-14 140 && \
+  update-alternatives --install /usr/bin/c++ c++ /usr/bin/g++-14 140 && \
+  apt-get clean && \
+  rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -s /bin/zsh -u 1000 -G sudo dev && \
+  echo 'dev ALL=(ALL) NOPASSWD:ALL' > /etc/sudoers.d/dev && \
+  chmod 0440 /etc/sudoers.d/dev && \
+  mkdir -p /workspace && \
+  chown dev:dev /workspace
+
+USER dev
+
+RUN git config --global --add safe.directory /workspace && \
+  mix local.hex --force && \
+  mix local.rebar --force
+
+WORKDIR /workspace
diff --git a/.devcontainer/devcontainer-lock.json b/.devcontainer/devcontainer-lock.json
new file mode 100644
index 0000000..523de18
--- /dev/null
+++ b/.devcontainer/devcontainer-lock.json
@@ -0,0 +1,24 @@
+{
+  "features": {
+    "ghcr.io/devcontainers/features/common-utils:2": {
+      "version": "2.5.9",
+      "resolved": "ghcr.io/devcontainers/features/common-utils@sha256:cb0c4d3c276f157eed17935747e364178d75fee17f55c4e129966f64633deb3a",
+      "integrity": "sha256:cb0c4d3c276f157eed17935747e364178d75fee17f55c4e129966f64633deb3a"
+    },
+    "ghcr.io/devcontainers/features/docker-outside-of-docker:1": {
+      "version": "1.10.0",
+      "resolved": "ghcr.io/devcontainers/features/docker-outside-of-docker@sha256:c2c2cf829505ead8e4892c88c31b6594ae94a2bbb209e16e1fac456c1a3a624e",
+      "integrity": "sha256:c2c2cf829505ead8e4892c88c31b6594ae94a2bbb209e16e1fac456c1a3a624e"
+    },
+    "ghcr.io/devcontainers/features/git:1": {
+      "version": "1.3.7",
+      "resolved": "ghcr.io/devcontainers/features/git@sha256:a3e43ff91b9f5f6bd2c14bd510d43e5e698f1266dc41027ba4e04e7e45be607a",
+      "integrity": "sha256:a3e43ff91b9f5f6bd2c14bd510d43e5e698f1266dc41027ba4e04e7e45be607a"
+    },
+    "ghcr.io/devcontainers/features/github-cli:1": {
+      "version": "1.1.0",
+      "resolved": "ghcr.io/devcontainers/features/github-cli@sha256:d22f50b70ed75339b4eed1ba9ecde3a1791f90e88d37936517e3bace0bbad671",
+      "integrity": "sha256:d22f50b70ed75339b4eed1ba9ecde3a1791f90e88d37936517e3bace0bbad671"
+    }
+  }
+}
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
new file mode 100644
index 0000000..c452d66
--- /dev/null
+++ b/.devcontainer/devcontainer.json
@@ -0,0 +1,33 @@
+{
+  "name": "EXGBoost Erlang 29 GCC 14",
+  "build": {
+    "dockerfile": "Dockerfile",
+    "context": ".."
+  },
+  "workspaceMount": "source=${localWorkspaceFolder},target=/workspace,type=bind,consistency=cached",
+  "workspaceFolder": "/workspace",
+  "containerEnv": {
+    "CC": "gcc-14",
+    "CXX": "g++-14",
+    "MIX_ENV": "dev"
+  },
+  "customizations": {
+    "vscode": {
+      "extensions": [
+        "jakebecker.elixir-ls"
+      ]
+    }
+  },
+  "postCreateCommand": "mix deps.get",
+  "remoteUser": "dev",
+  "features": {
+    "ghcr.io/devcontainers/features/common-utils:2": {
+      "configureZshAsDefaultShell": true
+    },
+    "ghcr.io/devcontainers/features/git:1": {},
+    "ghcr.io/devcontainers/features/github-cli:1": {},
+    "ghcr.io/devcontainers/features/docker-outside-of-docker:1": {
+      "moby": false
+    }
+  }
+}
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..5ae3052
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,12 @@
+version: 2
+updates:
+  - package-ecosystem: "mix"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 10
+
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
\ No newline at end of file
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..1c0bf8d
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,53 @@
+name: ci
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  test:
+    name: Test OTP ${{ matrix.otp }} / Elixir ${{ matrix.elixir }}
+    runs-on: ubuntu-24.04
+    env:
+      CC: gcc-14
+      CXX: g++-14
+      MIX_ENV: test
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - otp: "28.3"
+            elixir: "1.18.4"
+          - otp: "28.3"
+            elixir: "1.19.5"
+          - otp: "29.0.2"
+            elixir: "1.20.2"
+
+    steps:
+      - uses: actions/checkout@v5
+
+      - name: Install system dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y build-essential cmake ninja-build gcc-14 g++-14
+
+      - uses: erlef/setup-beam@v1
+        with:
+          otp-version: ${{ matrix.otp }}
+          elixir-version: ${{ matrix.elixir }}
+
+      - name: Install Hex and Rebar
+        run: |
+          mix local.hex --force
+          mix local.rebar --force
+
+      - name: Install dependencies
+        run: mix deps.get
+
+      - name: Check formatting
+        run: mix format --check-formatted
+
+      - name: Run tests
+        run: MIX_ENV=test mix test
diff --git a/.github/workflows/precompile.yml b/.github/workflows/precompile.yml
index dfbae0c..1494c92 100644
--- a/.github/workflows/precompile.yml
+++ b/.github/workflows/precompile.yml
@@ -10,20 +10,32 @@ jobs:
     runs-on: ubuntu-24.04
     env:
       MIX_ENV: "prod"
+      CC: gcc-14
+      CXX: g++-14
     strategy:
       matrix:
-        otp: ["27.0", "28.0"]
-        elixir: ["1.18.4"]
+        include:
+          - otp: "28.3"
+            elixir: "1.18.4"
+          - otp: "28.3"
+            elixir: "1.19.5"
+          - otp: "29.0.2"
+            elixir: "1.20.2"
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v5
       - uses: erlef/setup-beam@v1
         with:
           otp-version: ${{matrix.otp}}
           elixir-version: ${{matrix.elixir}}
-      - name: Install system dependecies
+      - name: Install Hex and Rebar
+        run: |
+          mix local.hex --force
+          mix local.rebar --force
+      - name: Install system dependencies
         run: |
           sudo apt-get update
           sudo apt-get install -y build-essential automake autoconf pkg-config bc m4 unzip zip \
+            cmake ninja-build gcc-14 g++-14 \
             gcc-aarch64-linux-gnu g++-aarch64-linux-gnu \
             gcc-riscv64-linux-gnu g++-riscv64-linux-gnu
       - name: Mix Test
@@ -42,25 +54,16 @@ jobs:
             cache/*.tar.gz
 
   macos:
-    runs-on: ${{matrix.runner}}
-    # Homebrew supports versioned Erlang/OTP but not Elixir
-    # It's a deliberate design decision from Homebrew to
-    # only support versioned distributions for certain packages
-    name: Mac (ARM) Erlang/OTP ${{matrix.otp}} / Elixir
+    runs-on: macos-15
+    name: Mac (ARM) Erlang/OTP 29.0.2 / Elixir 1.20.2
     env:
       MIX_ENV: "prod"
-    strategy:
-      matrix:
-        runner: ["macos-15"]
-        otp: ["27.0", "28.0"]
-        elixir: ["1.18.4"]
     steps:
-      - uses: actions/checkout@v3
-      - uses: asdf-vm/actions/install@v2
+      - uses: actions/checkout@v5
+      - uses: erlef/setup-beam@v1
         with:
-          tool_versions: |
-            erlang ${{matrix.otp}}
-            elixir ${{matrix.elixir}}
+          otp-version: "29.0.2"
+          elixir-version: "1.20.2"
       - name: Install libomp
         run: |
           brew install libomp
diff --git a/.gitignore b/.gitignore
index 15fd0cd..839a055 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,7 +9,6 @@ erl_crash.dump
 *.beam
 /config/*.secret.exs
 .elixir_ls/
-.tool-versions
 .vscode/
 checksum.exs
 .DS_Store
\ No newline at end of file
diff --git a/Makefile b/Makefile
index 1768938..2fd89e0 100644
--- a/Makefile
+++ b/Makefile
@@ -1,17 +1,28 @@
 # Environment variables passed via elixir_make
 # ERTS_INCLUDE_DIR
 # MIX_APP_PATH
+# MIX_ENV
 
 TEMP ?= $(HOME)/.cache
+MIX_ENV ?= dev
 XGBOOST_CACHE ?= $(TEMP)/exgboost
 XGBOOST_GIT_REPO ?= https://github.com/dmlc/xgboost.git
-# v3.0.5 tagged release
-XGBOOST_GIT_REV ?= v3.0.5
+# v3.1.3 tagged release
+XGBOOST_GIT_REV ?= v3.1.3
+OLD_XGBOOST_GIT_REV ?= v3.0.5
+NEW_XGBOOST_GIT_REV ?= $(XGBOOST_GIT_REV)
 XGBOOST_NS = xgboost-$(XGBOOST_GIT_REV)
 XGBOOST_DIR = $(XGBOOST_CACHE)/$(XGBOOST_NS)
 XGBOOST_LIB_DIR = $(XGBOOST_DIR)/build/xgboost
 XGBOOST_LIB_DIR_FLAG = $(XGBOOST_LIB_DIR)/exgboost.ok
 
+# Set build type based on MIX_ENV
+ifeq ($(MIX_ENV), prod)
+	CMAKE_BUILD_TYPE = Release
+else
+	CMAKE_BUILD_TYPE = RelWithDebInfo
+endif
+
 # Private configuration
 PRIV_DIR = $(MIX_APP_PATH)/priv
 EXGBOOST_DIR = $(realpath c/exgboost)
@@ -43,21 +54,21 @@ else
 endif
 
 $(EXGBOOST_SO): $(EXGBOOST_CACHE_SO)
-	@ mkdir -p $(PRIV_DIR)
+	@mkdir -p $(PRIV_DIR)
 	cp -a $(abspath $(EXGBOOST_CACHE_LIB_DIR)) $(EXGBOOST_LIB_DIR) ; \
 	cp -a $(abspath $(EXGBOOST_CACHE_SO)) $(EXGBOOST_SO) ;
 
 $(EXGBOOST_CACHE_SO): $(XGBOOST_LIB_DIR_FLAG) $(C_SRCS)
 	@mkdir -p cache
-	cp -a $(XGBOOST_LIB_DIR) $(EXGBOOST_CACHE_LIB_DIR)
-	mv $(XGBOOST_LIB_DIR)/lib/$(LIBXGBOOST) $(EXGBOOST_CACHE_LIB_DIR)
+	cp -R $(XGBOOST_LIB_DIR) $(EXGBOOST_CACHE_LIB_DIR)
+	cp $(XGBOOST_DIR)/lib/$(LIBXGBOOST) $(EXGBOOST_CACHE_LIB_DIR)
 	$(CC) $(CFLAGS) $(wildcard $(EXGBOOST_DIR)/src/*.c) $(LDFLAGS) -o $(EXGBOOST_CACHE_SO)
 	$(POST_INSTALL)
 
 # This new target handles fetching the source code.
 # It only runs if the .git directory inside the source folder is missing.
 $(XGBOOST_DIR)/.git:
-	mkdir -p $(XGBOOST_DIR) && \
+	@mkdir -p $(XGBOOST_DIR) && \
 		cd $(XGBOOST_DIR) && \
 		git init && \
 		git remote add origin $(XGBOOST_GIT_REPO) && \
@@ -69,10 +80,30 @@ $(XGBOOST_DIR)/.git:
 # It only contains the build commands.
 $(XGBOOST_LIB_DIR_FLAG): $(XGBOOST_DIR)/.git
 	cd $(XGBOOST_DIR) && \
-		cmake -B build -S . -DCMAKE_INSTALL_PREFIX=$(XGBOOST_LIB_DIR) -DCMAKE_BUILD_TYPE=RelWithDebInfo -GNinja $(CMAKE_FLAGS) && \
+		cmake -B build -S . -DCMAKE_INSTALL_PREFIX=$(XGBOOST_LIB_DIR) -DCMAKE_BUILD_TYPE=$(CMAKE_BUILD_TYPE) -GNinja $(CMAKE_FLAGS) && \
 		ninja -C build install
 	touch $(XGBOOST_LIB_DIR_FLAG)
 
+check-xgboost-c-api: $(XGBOOST_LIB_DIR_FLAG)
+	./scripts/check_xgboost_c_api.sh "$(XGBOOST_LIB_DIR)/include"
+
+compare-xgboost-c-api:
+	@set -eu; \
+	for rev in "$(OLD_XGBOOST_GIT_REV)" "$(NEW_XGBOOST_GIT_REV)"; do \
+		dir="$(XGBOOST_CACHE)/xgboost-$$rev"; \
+		mkdir -p "$$dir"; \
+		if [ ! -d "$$dir/.git" ]; then \
+			git -C "$$dir" init; \
+			git -C "$$dir" remote add origin "$(XGBOOST_GIT_REPO)"; \
+		fi; \
+		git -C "$$dir" fetch --depth 1 --recurse-submodules origin "$$rev"; \
+		git -C "$$dir" checkout -f FETCH_HEAD; \
+		git -C "$$dir" submodule update --init --recursive; \
+	done; \
+	./scripts/check_xgboost_c_api.sh --compare \
+		"$(XGBOOST_CACHE)/xgboost-$(OLD_XGBOOST_GIT_REV)/include" \
+		"$(XGBOOST_CACHE)/xgboost-$(NEW_XGBOOST_GIT_REV)/include"
+
 clean:
 	rm -rf $(EXGBOOST_CACHE_SO)
 	rm -rf $(EXGBOOST_CACHE_LIB_DIR)
diff --git a/README.md b/README.md
index e11c5f1..be91165 100644
--- a/README.md
+++ b/README.md
@@ -59,7 +59,7 @@ EXGBoost.predict(model, x)
 
 EXGBoost is designed to feel familiar to the users of the Python XGBoost library. `EXGBoost.train/2` is the
 primary entry point for training a model. It accepts a Nx tensor for the features and a Nx tensor for the labels.
-`EXGBoost.train/2` returns a trained`Booster` struct that can be used for prediction. `EXGBoost.train/2` also
+`EXGBoost.train/2` returns a trained `Booster` struct that can be used for prediction. `EXGBoost.train/2` also
 accepts a keyword list of options that can be used to configure the training process. See the
 [XGBoost documentation](https://xgboost.readthedocs.io/en/latest/parameter.html) for the full list of options.
 
@@ -119,6 +119,18 @@ It accepts a `Booster` struct (which is the output of `EXGBoost.train/2`).
 preds = EXGBoost.train(X, y) |> EXGBoost.predict(X)
 ```
 
+## Concurrency and Thread Safety
+
+**Important**: Booster objects are **not thread-safe** for concurrent predictions. The underlying XGBoost C API does not provide synchronization mechanisms, and sharing a single booster reference across multiple Elixir processes for concurrent predictions can lead to race conditions, memory corruption, or incorrect results. For this reason it is not recommended that you cache boosters to be used
+by multiple tasks in calling applications.
+
+### Why This Matters
+
+- `EXGBoost.predict/2` and `EXGBoost.inplace_predict/2` both use dirty CPU-bound NIF schedulers
+- This prevents blocking the BEAM scheduler but **does not** provide thread safety
+- Concurrent access to the same booster from multiple processes can cause undefined behavior
+
+
 ## Serialization
 
   A Booster can be serialized to a file using `EXGBoost.write_*` and loaded from a file
@@ -175,7 +187,7 @@ preds = EXGBoost.train(X, y) |> EXGBoost.predict(X)
 
 ### Precompiled Distribution
 
-We currenly offer the following precompiled packages for EXGBoost:
+We currently offer the following precompiled packages for EXGBoost:
 
 ```elixir
 %{
@@ -220,3 +232,22 @@ You also need to set `CC_PRECOMPILER_PRECOMPILE_ONLY_LOCAL=true` before the firs
 ## License
 
 Licensed under an [Apache-2](https://github.com/acalejos/exgboost/blob/main/LICENSE) license.
+
+## Maintaining this fork
+
+Recommended fork maintenance cadence:
+
+- Keep dependencies current with Dependabot PRs (configured in `.github/dependabot.yml`).
+- Ensure all PRs pass CI (`.github/workflows/ci.yml`) before merge.
+- Keep precompiled artefacts and checksums aligned with each release tag (`.github/workflows/precompile.yml`).
+- When upgrading `nx`, verify supported Elixir/OTP versions in CI matrix and `mix.exs`.
+
+Native compatibility checks:
+
+- Run `make check-xgboost-c-api` after changing `XGBOOST_GIT_REV` or native C files.
+- Run `mix test test/nif_test.exs` to validate runtime behavior after the API check passes.
+- To compare API declarations between two XGBoost versions, run:
+  `scripts/check_xgboost_c_api.sh --compare <old-include-dir> <new-include-dir>`
+- Or use make to fetch and compare tags directly:
+  `make compare-xgboost-c-api OLD_XGBOOST_GIT_REV=v3.0.5 NEW_XGBOOST_GIT_REV=v3.1.3`
+
diff --git a/c/exgboost/include/utils.h b/c/exgboost/include/utils.h
index a995657..8a88bbb 100644
--- a/c/exgboost/include/utils.h
+++ b/c/exgboost/include/utils.h
@@ -42,4 +42,8 @@ int exg_get_string_list(ErlNifEnv *env, ERL_NIF_TERM term, char ***out,
 int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term,
                          DMatrixHandle **dmats, unsigned *len);
 
+void exg_free_string_list(char **items, unsigned len);
+
+void exg_free_dmatrix_list(DMatrixHandle *dmats);
+
 #endif
\ No newline at end of file
diff --git a/c/exgboost/src/booster.c b/c/exgboost/src/booster.c
index 7f578f1..b0a23a2 100644
--- a/c/exgboost/src/booster.c
+++ b/c/exgboost/src/booster.c
@@ -7,6 +7,7 @@ static ERL_NIF_TERM make_Booster_resource(ErlNifEnv *env,
       enif_alloc_resource(Booster_RESOURCE_TYPE, sizeof(BoosterHandle *));
   if (resource != NULL) {
     *resource = handle;
+    // BEAM resource now owns the handle and releases it in resource cleanup.
     ret = exg_ok(env, enif_make_resource(env, resource));
     enif_release_resource(resource);
   } else {
@@ -49,6 +50,8 @@ ERL_NIF_TERM EXGBoosterCreate(ErlNifEnv *env, int argc,
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  // exg_get_dmatrix_list allocates this temporary array.
+  exg_free_dmatrix_list(dmats);
   return ret;
 }
 
@@ -148,6 +151,7 @@ ERL_NIF_TERM EXGBoosterSetParam(ErlNifEnv *env, int argc,
     ret = exg_error(env, "Booster parameter value must be a string");
     goto END;
   }
+  // XGBoost consumes name/value during this call; no ownership transfer.
   result = XGBoosterSetParam(booster, name, value);
   if (result == 0) {
     ret = enif_make_atom(env, "ok");
@@ -155,6 +159,12 @@ ERL_NIF_TERM EXGBoosterSetParam(ErlNifEnv *env, int argc,
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  if (name != NULL) {
+    enif_free(name);
+  }
+  if (value != NULL) {
+    enif_free(value);
+  }
   return ret;
 }
 
@@ -290,7 +300,7 @@ ERL_NIF_TERM EXGBoosterEvalOneIter(ErlNifEnv *env, int argc,
   int iter = -1;
   unsigned num_dmats = 0;
   unsigned num_evnames = 0;
-  char *out = NULL;
+  const char *out = NULL;
   ERL_NIF_TERM ret = -1;
   int result = -1;
   if (4 != argc) {
@@ -319,7 +329,7 @@ ERL_NIF_TERM EXGBoosterEvalOneIter(ErlNifEnv *env, int argc,
     ret = exg_error(env, "dmats and evnames must have the same length");
     goto END;
   }
-  result = XGBoosterEvalOneIter(booster, iter, dmats, evnames,
+  result = XGBoosterEvalOneIter(booster, iter, dmats, (const char **)evnames,
                                 (bst_ulong)num_dmats, &out);
   if (result == 0) {
     ret = exg_ok(env, enif_make_string(env, out, ERL_NIF_LATIN1));
@@ -327,6 +337,9 @@ ERL_NIF_TERM EXGBoosterEvalOneIter(ErlNifEnv *env, int argc,
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  // Helper-allocated arrays must be reclaimed on all paths.
+  exg_free_dmatrix_list(dmats);
+  exg_free_string_list(evnames, num_evnames);
   return ret;
 }
 
@@ -335,7 +348,7 @@ ERL_NIF_TERM EXGBoosterGetAttr(ErlNifEnv *env, int argc,
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
   char *key = NULL;
-  char *out = NULL;
+  const char *out = NULL;
   ERL_NIF_TERM ret = -1;
   int result = -1;
   int success = -1;
@@ -364,6 +377,9 @@ ERL_NIF_TERM EXGBoosterGetAttr(ErlNifEnv *env, int argc,
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  if (key != NULL) {
+    enif_free(key);
+  }
   return ret;
 }
 
@@ -415,6 +431,12 @@ ERL_NIF_TERM EXGBoosterSetAttr(ErlNifEnv *env, int argc,
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  if (key != NULL) {
+    enif_free(key);
+  }
+  if (value != NULL) {
+    enif_free(value);
+  }
   return ret;
 }
 
@@ -422,7 +444,7 @@ ERL_NIF_TERM EXGBoosterGetAttrNames(ErlNifEnv *env, int argc,
                                     const ERL_NIF_TERM argv[]) {
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
-  char **out = NULL;
+  const char **out = NULL;
   bst_ulong out_len = 0;
   ERL_NIF_TERM ret = -1;
   int result = -1;
@@ -440,10 +462,7 @@ ERL_NIF_TERM EXGBoosterGetAttrNames(ErlNifEnv *env, int argc,
   if (result == 0) {
     ERL_NIF_TERM arr[out_len];
     for (bst_ulong i = 0; i < out_len; ++i) {
-      char *local = enif_alloc(strlen(out[i]) + 1);
-      strcpy(local, out[i]);
-      arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1);
-      // TODO: Do we free here or is it handled by the XGBoost library / BEAM?
+      arr[i] = enif_make_string(env, out[i], ERL_NIF_LATIN1);
     }
     ret = exg_ok(env, enif_make_list_from_array(env, arr, out_len));
   } else {
@@ -485,7 +504,8 @@ ERL_NIF_TERM EXGBoosterSetStrFeatureInfo(ErlNifEnv *env, int argc,
     goto END;
   }
   handle = *resource;
-  result = XGBoosterSetStrFeatureInfo(handle, field, features, num_features);
+  result = XGBoosterSetStrFeatureInfo(handle, field, (const char **)features,
+                                      num_features);
   if (result == 0) {
     ret = ok_atom(env);
   } else {
@@ -493,8 +513,10 @@ ERL_NIF_TERM EXGBoosterSetStrFeatureInfo(ErlNifEnv *env, int argc,
   }
 END:
   if (features != NULL) {
-    enif_free(features);
-    features = NULL;
+    exg_free_string_list(features, num_features);
+  }
+  if (field != NULL) {
+    enif_free(field);
   }
   return ret;
 }
@@ -531,10 +553,8 @@ ERL_NIF_TERM EXGBoosterGetStrFeatureInfo(ErlNifEnv *env, int argc,
   if (result == 0) {
     ERL_NIF_TERM arr[out_size];
     for (bst_ulong i = 0; i < out_size; ++i) {
-      char *local = enif_alloc(strlen(c_out_features[i]) + 1);
-      strcpy(local, c_out_features[i]);
-      arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1);
-      // TODO: Do we free here or is it handled by the XGBoost library / BEAM?
+      // enif_make_string materializes a BEAM term; no temporary C copy needed.
+      arr[i] = enif_make_string(env, c_out_features[i], ERL_NIF_LATIN1);
     }
     ret = exg_ok(env, enif_make_list_from_array(env, arr, out_size));
   } else {
@@ -551,12 +571,12 @@ ERL_NIF_TERM EXGBoosterFeatureScore(ErlNifEnv *env, int argc,
                                     const ERL_NIF_TERM argv[]) {
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
-  char **config = NULL;
+  char *config = NULL;
   bst_ulong out_n_features = 0;
-  char **out_features = NULL;
+  const char **out_features = NULL;
   bst_ulong out_dim = 0;
-  bst_ulong *out_shape = NULL;
-  float *out_scores = NULL;
+  const bst_ulong *out_shape = NULL;
+  const float *out_scores = NULL;
   ERL_NIF_TERM ret = -1;
   int result = -1;
   if (2 != argc) {
@@ -596,13 +616,16 @@ ERL_NIF_TERM EXGBoosterFeatureScore(ErlNifEnv *env, int argc,
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  if (config != NULL) {
+    enif_free(config);
+  }
   return ret;
 }
 
 static ERL_NIF_TERM collect_prediction_results(ErlNifEnv *env,
-                                               bst_ulong *out_shape,
+                                               const bst_ulong *out_shape,
                                                bst_ulong out_dim,
-                                               float *out_result) {
+                                               const float *out_result) {
   bst_ulong out_len = 1;
   ERL_NIF_TERM shape_arr[out_dim];
   for (bst_ulong j = 0; j < out_dim; ++j) {
@@ -612,6 +635,7 @@ static ERL_NIF_TERM collect_prediction_results(ErlNifEnv *env,
   ERL_NIF_TERM shape = enif_make_tuple_from_array(env, shape_arr, out_dim);
   ERL_NIF_TERM result_arr[out_len];
   for (bst_ulong i = 0; i < out_len; ++i) {
+    // Values are copied into BEAM-managed terms here.
     result_arr[i] = enif_make_double(env, out_result[i]);
   }
   return exg_ok(env, enif_make_tuple2(
@@ -626,9 +650,9 @@ ERL_NIF_TERM EXGBoosterPredictFromDMatrix(ErlNifEnv *env, int argc,
   DMatrixHandle dmatrix;
   DMatrixHandle **dmatrix_resource = NULL;
   char *config = NULL;
-  bst_ulong *out_shape = NULL;
+  const bst_ulong *out_shape = NULL;
   bst_ulong out_dim = 0;
-  float *out_result = NULL;
+  const float *out_result = NULL;
 
   ERL_NIF_TERM ret = -1;
   int result = -1;
@@ -674,9 +698,9 @@ ERL_NIF_TERM EXGBoosterPredictFromDense(ErlNifEnv *env, int argc,
   DMatrixHandle **proxy_resource = NULL;
   char *values = NULL;
   char *config = NULL;
-  bst_ulong *out_shape = NULL;
+  const bst_ulong *out_shape = NULL;
   bst_ulong out_dim = 0;
-  float *out_result = NULL;
+  const float *out_result = NULL;
   int result = -1;
   ERL_NIF_TERM ret = -1;
   if (4 != argc) {
@@ -729,10 +753,10 @@ ERL_NIF_TERM EXGBoosterPredictFromCSR(ErlNifEnv *env, int argc,
   char *indices = NULL;
   char *data = NULL;
   char *config = NULL;
-  bst_ulong ncols = 0;
-  bst_ulong *out_shape = NULL;
+  int ncols = 0;
+  const bst_ulong *out_shape = NULL;
   bst_ulong out_dim = 0;
-  float *out_result = NULL;
+  const float *out_result = NULL;
   int result = -1;
   ERL_NIF_TERM ret = -1;
   if (7 != argc) {
@@ -760,6 +784,10 @@ ERL_NIF_TERM EXGBoosterPredictFromCSR(ErlNifEnv *env, int argc,
     ret = exg_error(env, "Ncols must be an integer");
     goto END;
   }
+  if (ncols < 0) {
+    ret = exg_error(env, "Ncols must be non-negative");
+    goto END;
+  }
   if (!exg_get_string(env, argv[5], &config)) {
     ret = exg_error(env, "Config must be a JSON-encoded string");
     goto END;
@@ -772,8 +800,9 @@ ERL_NIF_TERM EXGBoosterPredictFromCSR(ErlNifEnv *env, int argc,
   }
   booster = *booster_resource;
   result =
-      XGBoosterPredictFromCSR(booster, indptr, indices, data, ncols, config,
-                              proxy, &out_shape, &out_dim, &out_result);
+      XGBoosterPredictFromCSR(booster, indptr, indices, data,
+                              (bst_ulong)ncols, config, proxy, &out_shape,
+                              &out_dim, &out_result);
   if (result == 0) {
     ret = collect_prediction_results(env, out_shape, out_dim, out_result);
   } else {
@@ -818,6 +847,7 @@ ERL_NIF_TERM EXGBoosterLoadModel(ErlNifEnv *env, int argc,
   if (result == 0) {
     ret = make_Booster_resource(env, booster);
   } else {
+    XGBoosterFree(booster);
     ret = exg_error(env, XGBGetLastError());
   }
 END:
@@ -866,7 +896,7 @@ ERL_NIF_TERM EXGBoosterSerializeToBuffer(ErlNifEnv *env, int argc,
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
   bst_ulong out_len = 0;
-  char *out_buf = NULL;
+  const char *out_buf = NULL;
   int result = -1;
   ERL_NIF_TERM ret = -1;
   ErlNifBinary out_bin;
@@ -920,6 +950,7 @@ ERL_NIF_TERM EXGBoosterDeserializeFromBuffer(ErlNifEnv *env, int argc,
   if (result == 0) {
     ret = make_Booster_resource(env, booster);
   } else {
+    XGBoosterFree(booster);
     ret = exg_error(env, XGBGetLastError());
   }
 END:
@@ -955,6 +986,7 @@ ERL_NIF_TERM EXGBoosterLoadModelFromBuffer(ErlNifEnv *env, int argc,
   if (result == 0) {
     ret = make_Booster_resource(env, booster);
   } else {
+    XGBoosterFree(booster);
     ret = exg_error(env, XGBGetLastError());
   }
 END:
@@ -969,7 +1001,7 @@ ERL_NIF_TERM EXGBoosterSaveModelToBuffer(ErlNifEnv *env, int argc,
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
   bst_ulong out_len = 0;
-  char *out_buf = NULL;
+  const char *out_buf = NULL;
   char *config = NULL;
   int result = -1;
   ERL_NIF_TERM ret = -1;
@@ -1012,7 +1044,7 @@ ERL_NIF_TERM EXGBoosterSaveJsonConfig(ErlNifEnv *env, int argc,
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
   bst_ulong out_len = 0;
-  char *out_buf = NULL;
+  const char *out_buf = NULL;
   int result = -1;
   ERL_NIF_TERM ret = -1;
   ErlNifBinary out_bin;
@@ -1083,7 +1115,7 @@ ERL_NIF_TERM EXGBoosterDumpModelEx(ErlNifEnv *env, int argc,
   BoosterHandle booster;
   BoosterHandle **booster_resource = NULL;
   bst_ulong out_len = 0;
-  char **out_dump_array = NULL;
+  const char **out_dump_array = NULL;
   char *fmap = NULL;
   int with_stats = 0;
   char *format = NULL;
@@ -1116,10 +1148,7 @@ ERL_NIF_TERM EXGBoosterDumpModelEx(ErlNifEnv *env, int argc,
   if (result == 0) {
     ERL_NIF_TERM arr[out_len];
     for (bst_ulong i = 0; i < out_len; ++i) {
-      char *local = enif_alloc(strlen(out_dump_array[i]) + 1);
-      strcpy(local, out_dump_array[i]);
-      arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1);
-      // TODO: Do we free here or is it handled by the XGBoost library / BEAM ?
+      arr[i] = enif_make_string(env, out_dump_array[i], ERL_NIF_LATIN1);
     }
     ret = exg_ok(env, enif_make_list_from_array(env, arr, out_len));
   } else {
diff --git a/c/exgboost/src/config.c b/c/exgboost/src/config.c
index 89f432e..9590d66 100644
--- a/c/exgboost/src/config.c
+++ b/c/exgboost/src/config.c
@@ -59,7 +59,7 @@ ERL_NIF_TERM EXGBSetGlobalConfig(ErlNifEnv *env, int argc,
 
 ERL_NIF_TERM EXGBGetGlobalConfig(ErlNifEnv *env, int argc,
                                  const ERL_NIF_TERM argv[]) {
-  char *out = NULL;
+  const char *out = NULL;
   int result = -1;
   ERL_NIF_TERM ret = 0;
   if (argc != 0) {
@@ -68,7 +68,7 @@ ERL_NIF_TERM EXGBGetGlobalConfig(ErlNifEnv *env, int argc,
   }
   // No need to free out, it's a pointer to a static string defined in the
   // xgboost config struct
-  result = XGBGetGlobalConfig((char const **)&out);
+  result = XGBGetGlobalConfig(&out);
   if (result == 0) {
     ret = exg_ok(env, enif_make_string(env, out, ERL_NIF_LATIN1));
   } else {
@@ -76,4 +76,4 @@ ERL_NIF_TERM EXGBGetGlobalConfig(ErlNifEnv *env, int argc,
   }
 END:
   return ret;
-}
\ No newline at end of file
+}
diff --git a/c/exgboost/src/dmatrix.c b/c/exgboost/src/dmatrix.c
index 75dc621..28fcb74 100644
--- a/c/exgboost/src/dmatrix.c
+++ b/c/exgboost/src/dmatrix.c
@@ -7,6 +7,7 @@ static ERL_NIF_TERM make_DMatrix_resource(ErlNifEnv *env,
       enif_alloc_resource(DMatrix_RESOURCE_TYPE, sizeof(DMatrixHandle *));
   if (resource != NULL) {
     *resource = handle;
+    // BEAM resource now owns the handle and releases it in resource cleanup.
     ret = exg_ok(env, enif_make_resource(env, resource));
     enif_release_resource(resource);
   } else {
@@ -212,71 +213,6 @@ ERL_NIF_TERM EXGDMatrixCreateFromSparse(ErlNifEnv *env, int argc,
   return ret;
 }
 
-ERL_NIF_TERM EXGDMatrixCreateFromCSREx(ErlNifEnv *env, int argc,
-                                       const ERL_NIF_TERM argv[]) {
-  ErlNifBinary indptr_bin;
-  ErlNifBinary indices_bin;
-  ErlNifBinary data_bin;
-  int result = -1;
-  ErlNifUInt64 *indptr = NULL;
-  uint32_t *indices = NULL;
-  float *data = NULL;
-  uint32_t nindptr = 0;
-  ErlNifUInt64 nelem = 0;
-  ErlNifUInt64 ncol = 0;
-  DMatrixHandle handle;
-  ERL_NIF_TERM ret = 0;
-  if (argc != 6) {
-    ret = exg_error(env, "Wrong number of arguments");
-    goto END;
-  }
-  if (!enif_inspect_binary(env, argv[0], &indptr_bin)) {
-    ret = exg_error(env, "Indptr must be a binary of uint64_t");
-    goto END;
-  }
-  if (!enif_inspect_binary(env, argv[1], &indices_bin)) {
-    ret = exg_error(env, "Indices must be a binary of uint64_t");
-    goto END;
-  }
-  if (!enif_inspect_binary(env, argv[2], &data_bin)) {
-    ret = exg_error(env, "Data must be a binary of uint64_t");
-    goto END;
-  }
-  if (!enif_get_uint(env, argv[3], &nindptr)) {
-    ret = exg_error(env, "Nindptr must be a uint64_t");
-    goto END;
-  }
-  if (!enif_get_uint64(env, argv[4], &nelem)) {
-    ret = exg_error(env, "Nelem must be a uint64_t");
-    goto END;
-  }
-  if (!enif_get_uint64(env, argv[5], &ncol)) {
-    ret = exg_error(env, "Ncol must be a uint64_t");
-    goto END;
-  }
-  indptr = (ErlNifUInt64 *)indptr_bin.data;
-  indices = (uint32_t *)indices_bin.data;
-  data = (float *)data_bin.data;
-  if (indptr_bin.size != nindptr * sizeof(ErlNifUInt64)) {
-    ret = exg_error(env, "Indptr size does not match nindptr");
-    goto END;
-  }
-  if (data_bin.size != nelem * sizeof(float)) {
-    ret = exg_error(env, "Data size does not match nelem");
-    goto END;
-  }
-  result = XGDMatrixCreateFromCSREx(indptr, indices, data, nindptr, nelem, ncol,
-                                    &handle);
-  if (result == 0) {
-    ret = exg_ok(env, enif_make_resource(env, handle));
-    enif_release_resource(handle);
-  } else {
-    ret = exg_error(env, XGBGetLastError());
-  }
-END:
-  return ret;
-}
-
 ERL_NIF_TERM EXGDMatrixCreateFromDense(ErlNifEnv *env, int argc,
                                        const ERL_NIF_TERM argv[]) {
   int result = -1;
@@ -343,16 +279,21 @@ ERL_NIF_TERM EXGDMatrixSetStrFeatureInfo(ErlNifEnv *env, int argc,
     goto END;
   }
   handle = *resource;
-  result = XGDMatrixSetStrFeatureInfo(handle, field, features, num_features);
+  // XGBoost reads features during the call; caller keeps ownership.
+  result = XGDMatrixSetStrFeatureInfo(handle, field,
+                                      (const char **)features, num_features);
   if (result == 0) {
     ret = ok_atom(env);
   } else {
     ret = exg_error(env, XGBGetLastError());
   }
 END:
+  // Helper-allocated buffers must be reclaimed on all paths.
   if (features != NULL) {
-    enif_free(features);
-    features = NULL;
+    exg_free_string_list(features, num_features);
+  }
+  if (field != NULL) {
+    enif_free(field);
   }
   return ret;
 }
@@ -390,10 +331,8 @@ ERL_NIF_TERM EXGDMatrixGetStrFeatureInfo(ErlNifEnv *env, int argc,
   if (result == 0) {
     ERL_NIF_TERM arr[out_size];
     for (bst_ulong i = 0; i < out_size; ++i) {
-      char *local = enif_alloc(strlen(c_out_features[i]) + 1);
-      strcpy(local, c_out_features[i]);
-      arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1);
-      // TODO: Do we free here or is it handled by the XGBoost library / BEAM?
+      // enif_make_string materializes a BEAM term; no temporary C copy needed.
+      arr[i] = enif_make_string(env, c_out_features[i], ERL_NIF_LATIN1);
     }
     ret = exg_ok(env, enif_make_list_from_array(env, arr, out_size));
   } else {
@@ -646,7 +585,7 @@ ERL_NIF_TERM EXGDMatrixGetFloatInfo(ErlNifEnv *env, int argc,
   DMatrixHandle handle;
   DMatrixHandle **resource = NULL;
   char *field = NULL;
-  float *out = NULL;
+  const float *out = NULL;
   bst_ulong len = 0;
   int result = -1;
   ERL_NIF_TERM ret = 0;
@@ -679,6 +618,7 @@ ERL_NIF_TERM EXGDMatrixGetFloatInfo(ErlNifEnv *env, int argc,
   if (result == 0) {
     arr = enif_alloc(sizeof(ERL_NIF_TERM) * len);
     for (int i = 0; i < len; i++) {
+      // Values are copied into BEAM-managed terms here.
       arr[i] = enif_make_double(env, out[i]);
     }
     ret = exg_ok(env, enif_make_list_from_array(env, arr, len));
@@ -700,7 +640,7 @@ ERL_NIF_TERM EXGDMatrixGetUIntInfo(ErlNifEnv *env, int argc,
   DMatrixHandle handle;
   DMatrixHandle **resource = NULL;
   char *field = NULL;
-  unsigned *out = NULL;
+  const unsigned *out = NULL;
   bst_ulong len = 0;
   int result = -1;
   ERL_NIF_TERM ret = 0;
@@ -728,6 +668,7 @@ ERL_NIF_TERM EXGDMatrixGetUIntInfo(ErlNifEnv *env, int argc,
   if (result == 0) {
     arr = enif_alloc(sizeof(ERL_NIF_TERM) * len);
     for (int i = 0; i < len; i++) {
+      // Values are copied into BEAM-managed terms here.
       arr[i] = enif_make_uint(env, out[i]);
     }
     ret = exg_ok(env, enif_make_list_from_array(env, arr, len));
@@ -816,6 +757,7 @@ ERL_NIF_TERM EXGDMatrixGetDataAsCSR(ErlNifEnv *env, int argc,
                       enif_make_list_from_array(env, indices, num_non_missing),
                       enif_make_list_from_array(env, data, num_non_missing)));
 END:
+  // Mixed allocators: enif_free for enif_alloc buffers, free for malloc buffers.
   if (config != NULL) {
     enif_free(config);
     config = NULL;
@@ -832,6 +774,15 @@ ERL_NIF_TERM EXGDMatrixGetDataAsCSR(ErlNifEnv *env, int argc,
     free(out_data);
     out_data = NULL;
   }
+  if (indptr != NULL) {
+    enif_free(indptr);
+  }
+  if (indices != NULL) {
+    enif_free(indices);
+  }
+  if (data != NULL) {
+    enif_free(data);
+  }
   return ret;
 };
 
diff --git a/c/exgboost/src/utils.c b/c/exgboost/src/utils.c
index 401def7..b98aa14 100644
--- a/c/exgboost/src/utils.c
+++ b/c/exgboost/src/utils.c
@@ -61,12 +61,14 @@ int exg_get_list(ErlNifEnv *env, ERL_NIF_TERM term, double **out) {
     return 0;
   }
   *out = (double *)enif_alloc(len * sizeof(double));
-  if (out == NULL) {
+  if (*out == NULL) {
     return 0;
   }
   while (enif_get_list_cell(env, term, &head, &tail)) {
     int ret = enif_get_double(env, head, &((*out)[i]));
     if (!ret) {
+      enif_free(*out);
+      *out = NULL;
       return 0;
     }
     term = tail;
@@ -89,6 +91,8 @@ int exg_get_string_list(ErlNifEnv *env, ERL_NIF_TERM term, char ***out,
   while (enif_get_list_cell(env, term, &head, &tail)) {
     int ret = exg_get_string(env, head, &((*out)[i]));
     if (!ret) {
+      exg_free_string_list(*out, i);
+      *out = NULL;
       return 0;
     }
     term = tail;
@@ -105,13 +109,15 @@ int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term,
     return 0;
   }
   *dmats = (DMatrixHandle *)enif_alloc(*len * sizeof(DMatrixHandle));
-  if (NULL == dmats) {
+  if (NULL == *dmats) {
     return 0;
   }
   while (enif_get_list_cell(env, term, &head, &tail)) {
     DMatrixHandle **resource = NULL;
     if (!enif_get_resource(env, head, DMatrix_RESOURCE_TYPE,
                            (void *)&(resource))) {
+      exg_free_dmatrix_list(*dmats);
+      *dmats = NULL;
       return 0;
     }
     memcpy(&((*dmats)[i]), resource, sizeof(DMatrixHandle));
@@ -121,6 +127,24 @@ int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term,
   return 1;
 }
 
+void exg_free_string_list(char **items, unsigned len) {
+  if (items == NULL) {
+    return;
+  }
+  for (unsigned i = 0; i < len; ++i) {
+    if (items[i] != NULL) {
+      enif_free(items[i]);
+    }
+  }
+  enif_free(items);
+}
+
+void exg_free_dmatrix_list(DMatrixHandle *dmats) {
+  if (dmats != NULL) {
+    enif_free(dmats);
+  }
+}
+
 ERL_NIF_TERM exg_get_binary_address(ErlNifEnv *env, int argc,
                                     const ERL_NIF_TERM argv[]) {
   ErlNifBinary bin;
diff --git a/lib/exgboost/array_interface.ex b/lib/exgboost/array_interface.ex
index ad2fc8b..7e35ba3 100644
--- a/lib/exgboost/array_interface.ex
+++ b/lib/exgboost/array_interface.ex
@@ -4,7 +4,7 @@ defmodule EXGBoost.ArrayInterface do
 
   @typedoc """
   The XGBoost C API uses and is moving towards mainly supporting the use of
-  JSON-Encoded NumPy ArrayyInterface format to pass data to and from the C API. This struct
+  JSON-Encoded NumPy ArrayInterface format to pass data to and from the C API. This struct
   is used to represent the ArrayInterface format.
 
   If you wish to use the EXGBoost.NIF library directly, this will be the desired format
@@ -17,10 +17,10 @@ defmodule EXGBoost.ArrayInterface do
   @type t :: %__MODULE__{
           typestr: String.t(),
           shape: tuple(),
-          address: pos_integer(),
+          address: non_neg_integer(),
           readonly: boolean(),
-          tensor: Nx.Tensor.t(),
-          binary: binary()
+          tensor: Nx.Tensor.t() | nil,
+          binary: binary() | nil
         }
 
   @enforce_keys [:typestr, :shape, :address, :readonly]
@@ -94,11 +94,41 @@ defmodule EXGBoost.ArrayInterface do
       {"data", [address, readonly]}, acc ->
         [{:address, address} | [{:readonly, readonly} | acc]]
 
+      {:data, [address, readonly]}, acc ->
+        [{:address, address} | [{:readonly, readonly} | acc]]
+
       {"shape", shape}, acc ->
         [{:shape, List.to_tuple(shape)} | acc]
 
-      {key, value}, acc ->
-        [{String.to_existing_atom(key), value} | acc]
+      {:shape, shape}, acc ->
+        [{:shape, List.to_tuple(shape)} | acc]
+
+      {"typestr", typestr}, acc ->
+        [{:typestr, typestr} | acc]
+
+      {:typestr, typestr}, acc ->
+        [{:typestr, typestr} | acc]
+
+      {"version", version}, acc ->
+        [{:version, version} | acc]
+
+      {:version, version}, acc ->
+        [{:version, version} | acc]
+
+      {"tensor", tensor}, acc ->
+        [{:tensor, tensor} | acc]
+
+      {:tensor, tensor}, acc ->
+        [{:tensor, tensor} | acc]
+
+      {"binary", binary}, acc ->
+        [{:binary, binary} | acc]
+
+      {:binary, binary}, acc ->
+        [{:binary, binary} | acc]
+
+      {_key, _value}, acc ->
+        acc
     end)
     |> then(&struct(__MODULE__, &1))
   end
@@ -107,7 +137,7 @@ defmodule EXGBoost.ArrayInterface do
   This function is used to convert Nx.Tensors to the ArrayInterface format.
 
   Example:
-    iex> EXGBoost.from_tensor(Nx.tensor([[1,2,3],[4,5,6]]))
+    iex> EXGBoost.ArrayInterface.from_tensor(Nx.tensor([[1,2,3],[4,5,6]]))
         #ArrayInterface<
         %{data: [4418559984, true], shape: [2, 3], typestr: "<i8", version: 3}
   """
@@ -128,6 +158,7 @@ defmodule EXGBoost.ArrayInterface do
       end
 
     binary = Nx.to_binary(tensor)
+
     tensor_addr =
       EXGBoost.NIF.get_binary_address(binary) |> EXGBoost.Internal.unwrap!()
 
@@ -144,12 +175,33 @@ defmodule EXGBoost.ArrayInterface do
   @spec get_tensor(EXGBoost.ArrayInterface.t()) :: Nx.Tensor.t()
   def get_tensor(%__MODULE__{tensor: nil} = arr_int) do
     num_items = arr_int.shape |> Tuple.to_list() |> Enum.product()
-    <<_endianess::utf8, char_code::binary-size(1), bytes::binary>> = arr_int.typestr
+    <<endianess::utf8, char_code::binary-size(1), bytes::binary>> = arr_int.typestr
+
+    if endianess not in [?<, ?|] do
+      raise ArgumentError,
+            "Unsupported endianness in typestr #{inspect(arr_int.typestr)}. " <>
+              "Expected little-endian ('<') or non-endian ('|')."
+    end
+
+    bit_width = String.to_integer(bytes) * 8
 
     nx_type =
       case char_code do
-        "i" -> {:s, String.to_integer(bytes) * 8}
-        other -> {String.to_existing_atom(other), String.to_integer(bytes) * 8}
+        "i" ->
+          {:s, bit_width}
+
+        "u" ->
+          {:u, bit_width}
+
+        "f" ->
+          {:f, bit_width}
+
+        "c" ->
+          {:c, bit_width}
+
+        other ->
+          raise ArgumentError,
+                "Unsupported typestr code #{inspect(other)} in #{inspect(arr_int.typestr)}"
       end
 
     tensor_bin =
diff --git a/lib/exgboost/booster.ex b/lib/exgboost/booster.ex
index 1724ee2..adbdac4 100644
--- a/lib/exgboost/booster.ex
+++ b/lib/exgboost/booster.ex
@@ -294,7 +294,7 @@ defmodule EXGBoost.Booster do
           EXGBoost.NIF.booster_load_json_config(booster.ref, config) |> Internal.unwrap!()
 
         :model ->
-          EXGBoost.NIF.booster_deserialize_from_buffer(source) |> Internal.unwrap!()
+          load_model_with_fallback(source)
 
         :weights ->
           EXGBoost.NIF.booster_load_model_from_buffer(source) |> Internal.unwrap!()
@@ -303,6 +303,16 @@ defmodule EXGBoost.Booster do
     struct(booster, ref: booster_ref)
   end
 
+  defp load_model_with_fallback(source) do
+    case EXGBoost.NIF.booster_load_model_from_buffer(source) do
+      {:error, _reason} ->
+        EXGBoost.NIF.booster_deserialize_from_buffer(source) |> Internal.unwrap!()
+
+      result ->
+        Internal.unwrap!(result)
+    end
+  end
+
   @doc """
   Get a formatted representation of the Booster's model.
 
diff --git a/lib/exgboost/internal.ex b/lib/exgboost/internal.ex
index dd4a3a1..0d24ba0 100644
--- a/lib/exgboost/internal.ex
+++ b/lib/exgboost/internal.ex
@@ -139,6 +139,17 @@ defmodule EXGBoost.Internal do
   end
 
   def unwrap!({:ok, val}), do: val
-  def unwrap!({:error, reason}), do: raise(reason)
+
+  def unwrap!({:error, reason}) when is_exception(reason) do
+    raise(reason)
+  end
+
+  def unwrap!({:error, reason}) when is_list(reason) do
+    message = if List.ascii_printable?(reason), do: List.to_string(reason), else: inspect(reason)
+    raise ArgumentError, message
+  end
+
+  def unwrap!({:error, reason}) when is_binary(reason), do: raise(ArgumentError, reason)
+  def unwrap!({:error, reason}), do: raise(ArgumentError, inspect(reason))
   def unwrap!(:ok), do: :ok
 end
diff --git a/lib/exgboost/nif.ex b/lib/exgboost/nif.ex
index bbed9e0..c5086ea 100644
--- a/lib/exgboost/nif.ex
+++ b/lib/exgboost/nif.ex
@@ -355,13 +355,14 @@ defmodule EXGBoost.NIF do
           :ok | {:error, String.t()}
   def booster_save_model(_handle, _path), do: :erlang.nif_error(:not_implemented)
 
-  @spec booster_serialize_to_buffer(booster_reference()) :: binary()
+  @spec booster_serialize_to_buffer(booster_reference()) :: {:ok, binary()} | {:error, String.t()}
   def booster_serialize_to_buffer(_handle), do: :erlang.nif_error(:not_implemented)
 
   @spec booster_deserialize_from_buffer(binary()) :: exgboost_return_type(booster_reference())
   def booster_deserialize_from_buffer(_buffer), do: :erlang.nif_error(:not_implemented)
 
-  @spec booster_save_model_to_buffer(booster_reference(), String.t()) :: binary()
+  @spec booster_save_model_to_buffer(booster_reference(), String.t()) ::
+          {:ok, binary()} | {:error, String.t()}
   def booster_save_model_to_buffer(_handle, _config), do: :erlang.nif_error(:not_implemented)
 
   @spec booster_load_model_from_buffer(binary()) :: exgboost_return_type(booster_reference())
@@ -370,7 +371,7 @@ defmodule EXGBoost.NIF do
   @spec booster_load_json_config(booster_reference(), String.t()) :: :ok | {:error, String.t()}
   def booster_load_json_config(_handle, _config), do: :erlang.nif_error(:not_implemented)
 
-  @spec booster_save_json_config(booster_reference()) :: binary()
+  @spec booster_save_json_config(booster_reference()) :: {:ok, binary()} | {:error, String.t()}
   def booster_save_json_config(_handle), do: :erlang.nif_error(:not_implemented)
 
   def booster_dump_model(_handle, _fmap, _with_stats, _format),
diff --git a/lib/exgboost/plotting.ex b/lib/exgboost/plotting.ex
index 50719fb..6ca6f18 100644
--- a/lib/exgboost/plotting.ex
+++ b/lib/exgboost/plotting.ex
@@ -392,11 +392,16 @@ defmodule EXGBoost.Plotting do
     ]
   end
 
-  HTTPoison.start()
+  :inets.start()
+  :ssl.start()
+
+  @schema (fn ->
+            {:ok, {{_, 200, _}, _, body}} =
+              :httpc.request(:get, {"https://vega.github.io/schema/vega/v5.json", []}, [], [])
+
+            body |> to_string() |> Jason.decode!() |> ExJsonSchema.Schema.resolve()
+          end).()
 
-  @schema HTTPoison.get!("https://vega.github.io/schema/vega/v5.json").body
-          |> Jason.decode!()
-          |> ExJsonSchema.Schema.resolve()
 
   @mark_text_doc "Accepts a keyword list of Vega `text` Mark properties. Reference [here](https://vega.github.io/vega/docs/marks/text/) for more details. Accepts either a string (expected to be valid Vega property names) or Elixir-styled atom. Note that keys are snake-cased instead of camel-case (e.g. Vega `fontSize` becomes `font_size`)"
   @mark_rect_doc "Accepts a keyword list of Vega `rect` Mark properties. Reference [here](https://vega.github.io/vega/docs/marks/rect/) for more details. Accepts either a string (expected to be valid Vega property names) or Elixir-styled atom. Note that keys are snake-cased instead of camel-case (e.g. Vega `fontSize` becomes `font_size`)"
@@ -834,7 +839,7 @@ defmodule EXGBoost.Plotting do
   @spec get_defaults() :: Keyword.t()
   def get_defaults(), do: @defaults
 
-  @spec get_styles() :: [{atom(), [style(), ...]}, ...]
+  @spec get_styles() :: EXGBoost.Plotting.Style.t()
   def get_styles(), do: @styles
 
   defp validate_spec(spec) do
diff --git a/lib/exgboost/training/callback.ex b/lib/exgboost/training/callback.ex
index b3e6287..ee1d28a 100644
--- a/lib/exgboost/training/callback.ex
+++ b/lib/exgboost/training/callback.ex
@@ -43,6 +43,14 @@ defmodule EXGBoost.Training.Callback do
 
   """
   alias EXGBoost.Training.State
+
+  @type t :: %__MODULE__{
+          event: event(),
+          fun: fun(),
+          name: atom(),
+          init_state: any()
+        }
+
   @enforce_keys [:event, :fun]
   defstruct [:event, :fun, :name, :init_state]
 
@@ -54,7 +62,7 @@ defmodule EXGBoost.Training.Callback do
   @doc """
   Factory for a new callback with an initial state.
   """
-  @spec new(event :: event(), fun :: fun(), name :: atom(), init_state :: any()) :: Callback.t()
+  @spec new(event :: event(), fun :: fun(), name :: atom(), init_state :: any()) :: t()
   def new(event, fun, name, init_state \\ %{})
       when event in @valid_events and is_function(fun, 1) and is_atom(name) and not is_nil(name) do
     %__MODULE__{event: event, fun: fun, name: name, init_state: init_state}
diff --git a/lib/exgboost/training/state.ex b/lib/exgboost/training/state.ex
index de7cd90..4741409 100644
--- a/lib/exgboost/training/state.ex
+++ b/lib/exgboost/training/state.ex
@@ -1,5 +1,15 @@
 defmodule EXGBoost.Training.State do
   @moduledoc false
+
+  @type t :: %__MODULE__{
+          booster: EXGBoost.Booster.t(),
+          iteration: non_neg_integer(),
+          max_iteration: -1 | non_neg_integer(),
+          meta_vars: map(),
+          metrics: map(),
+          status: :cont | :halt
+        }
+
   @enforce_keys [:booster]
   defstruct [
     :booster,
diff --git a/mix.exs b/mix.exs
index 9a12ab9..e700be3 100644
--- a/mix.exs
+++ b/mix.exs
@@ -1,6 +1,7 @@
 defmodule EXGBoost.MixProject do
   use Mix.Project
-  @version "0.5.2"
+
+  @version "0.8.1"
 
   def project do
     [
@@ -8,34 +9,34 @@ defmodule EXGBoost.MixProject do
       version: @version,
       make_precompiler: {:nif, CCPrecompiler},
       make_precompiler_url:
-        "https://github.com/acalejos/exgboost/releases/download/v#{@version}/@{artefact_filename}",
+        "https://github.com/iperks/exgboost/releases/download/v#{@version}/@{artefact_filename}",
       make_precompiler_priv_paths: ["libexgboost.*", "lib"],
       # NIF Versions correspond to OTP Releases
       # https://github.com/erlang/otp/blob/d3aa6c044c3927f011fb76ac087d5ce0e814954c/erts/emulator/beam/erl_nif.h#L57
       make_precompiler_nif_versions: [
-        versions: ["2.15", "2.16", "2.17"]
+        versions: ["2.15", "2.16", "2.17", "2.18"]
       ],
-      elixir: "~> 1.14",
+      elixir: "~> 1.17",
       start_permanent: Mix.env() == :prod,
       compilers: [:elixir_make] ++ Mix.compilers(),
       deps: deps(),
       name: "EXGBoost",
-      source_url: "https://github.com/acalejos/exgboost",
-      homepage_url: "https://github.com/acalejos/exgboost",
+      source_url: "https://github.com/iperks/exgboost",
+      homepage_url: "https://github.com/iperks/exgboost",
       docs: docs(),
       package: package(),
-      preferred_cli_env: [
-        docs: :docs,
-        "hex.publish": :docs
-      ],
       before_closing_body_tag: &before_closing_body_tag/1,
       name: "EXGBoost",
       description:
         "Elixir bindings for the XGBoost library. `EXGBoost` provides an implementation of XGBoost that works with
-      [Nx](https://hexdocs.pm/nx/Nx.html) tensors."
+      [Nx](https://hexdocs.pm/nx/Nx.html) tensors. Maintained fork of acalejos/exgboost."
     ]
   end
 
+  def cli do
+    [preferred_envs: [docs: :docs, "hex.publish": :docs]]
+  end
+
   def application do
     [
       extra_applications: [:logger],
@@ -47,26 +48,26 @@ defmodule EXGBoost.MixProject do
     [
       {:elixir_make, "~> 0.4", runtime: false},
       {:nimble_options, "~> 1.0"},
-      {:nx, "~> 0.7"},
+      {:nx, "~> 0.9"},
       {:jason, "~> 1.3"},
-      {:ex_doc, "~> 0.31.0", only: :docs},
+      {:ex_doc, "~> 0.40", only: :docs},
       {:cc_precompiler, "~> 0.1.0", runtime: false},
       {:exterval, "0.2.0"},
-      {:ex_json_schema, "~> 0.11.0"},
-      {:httpoison, "~> 2.0", runtime: false},
+      {:ex_json_schema, "~> 0.11.4"},
       {:vega_lite, "~> 0.1"},
       {:vega_lite_convert, "~> 1.0.1"},
       {:kino, "~> 0.11"},
       {:scidata, "~> 0.1", only: :dev},
-      {:kino_vega_lite, "~> 0.1.9", only: :dev}
+      {:kino_vega_lite, "~> 0.1.9", only: :dev},
+      {:dialyxir, "~> 1.4", only: [:dev, :test], runtime: false}
     ]
   end
 
   defp package do
     [
-      maintainers: ["Andres Alejos"],
+      maintainers: ["Ian Perks"],
       licenses: ["Apache-2.0"],
-      links: %{"GitHub" => "https://github.com/acalejos/exgboost"},
+      links: %{"GitHub" => "https://github.com/iperks/exgboost"},
       files: [
         "lib",
         "mix.exs",
diff --git a/mix.lock b/mix.lock
index 452af39..4fa6628 100644
--- a/mix.lock
+++ b/mix.lock
@@ -1,37 +1,29 @@
 %{
   "castore": {:hex, :castore, "0.1.22", "4127549e411bedd012ca3a308dede574f43819fe9394254ca55ab4895abfa1a2", [:mix], [], "hexpm", "c17576df47eb5aa1ee40cc4134316a99f5cad3e215d5c77b8dd3cfef12a22cac"},
   "cc_precompiler": {:hex, :cc_precompiler, "0.1.11", "8c844d0b9fb98a3edea067f94f616b3f6b29b959b6b3bf25fee94ffe34364768", [:mix], [{:elixir_make, "~> 0.7", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "3427232caf0835f94680e5bcf082408a70b48ad68a5f5c0b02a3bea9f3a075b9"},
-  "certifi": {:hex, :certifi, "2.15.0", "0e6e882fcdaaa0a5a9f2b3db55b1394dba07e8d6d9bcad08318fb604c6839712", [:rebar3], [], "hexpm", "b147ed22ce71d72eafdad94f055165c1c182f61a2ff49df28bcc71d1d5b94a60"},
-  "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"},
-  "decimal": {:hex, :decimal, "2.3.0", "3ad6255aa77b4a3c4f818171b12d237500e63525c2fd056699967a3e7ea20f62", [:mix], [], "hexpm", "a4d66355cb29cb47c3cf30e71329e58361cfcb37c34235ef3bf1d7bf3773aeac"},
+  "complex": {:hex, :complex, "0.7.0", "695632ef9487517aa5d57edd1697801079d622414cb2e1a7cf538b1f9a50f205", [:mix], [], "hexpm", "0ee39c0803129f546e7f3f640da8f021c9e659402bf59da6f7f2c4848f068f8d"},
+  "decimal": {:hex, :decimal, "3.1.1", "430d87b04011ce6cbd4fd205be758311a81f87d552d40904abd00f015935b1d0", [:mix], [], "hexpm", "c5f25f2ced74a0587d03e6023f595db8e924c9d3922c8c8ffd9edfc4498cf1f6"},
+  "dialyxir": {:hex, :dialyxir, "1.4.7", "dda948fcee52962e4b6c5b4b16b2d8fa7d50d8645bbae8b8685c3f9ecb7f5f4d", [:mix], [{:erlex, ">= 0.2.8", [hex: :erlex, repo: "hexpm", optional: false]}], "hexpm", "b34527202e6eb8cee198efec110996c25c5898f43a4094df157f8d28f27d9efe"},
   "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"},
-  "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"},
-  "ex_doc": {:hex, :ex_doc, "0.31.2", "8b06d0a5ac69e1a54df35519c951f1f44a7b7ca9a5bb7a260cd8a174d6322ece", [:mix], [{:earmark_parser, "~> 1.4.39", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.1", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "317346c14febaba9ca40fd97b5b5919f7751fb85d399cc8e7e8872049f37e0af"},
-  "ex_json_schema": {:hex, :ex_json_schema, "0.11.1", "b593f92937a095f66054bb318681397dfe7304e7d2b6b1a7534ea3aa40024f8c", [:mix], [{:decimal, "~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}], "hexpm", "32d651a575a6ce2fd613f140b0fef8dd0acc7cf8e8bcd29a3a1be5c945700dd5"},
+  "elixir_make": {:hex, :elixir_make, "0.10.0", "16577e2583a79bb79237bbff349619ef5d80afffc07eac6e4faf0d00e2ddaf7d", [:mix], [], "hexpm", "dc1f09fb7fa68866b886abd5f0f3c83553b1a19a52359a899e92af1bb3b31982"},
+  "erlex": {:hex, :erlex, "0.2.8", "cd8116f20f3c0afe376d1e8d1f0ae2452337729f68be016ea544a72f767d9c12", [:mix], [], "hexpm", "9d66ff9fedf69e49dc3fd12831e12a8a37b76f8651dd21cd45fcf5561a8a7590"},
+  "ex_doc": {:hex, :ex_doc, "0.40.3", "4a972ffe64bc07dc605af487e98fc19b72a4185f55ca031b94c0552d6071c1d9", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "2756e357742fecd9749b489b85d67c9ce99c465f2e75728d9e6dc8d704b973de"},
+  "ex_json_schema": {:hex, :ex_json_schema, "0.11.4", "d2f7d31894d048f79ed6c5a76515c266d5bd137438c53fa39c55f6ae98a05f47", [:mix], [{:decimal, "~> 3.0", [hex: :decimal, repo: "hexpm", optional: false]}], "hexpm", "0bbe87044ef0154be2a91ab6927d69c5fcccdb21908a135653fc10dcbbb79c3b"},
   "exterval": {:hex, :exterval, "0.2.0", "cb2ce1dcaa08f7827d5eab25e0268e786b6b05b9b6d27f79d1d77f7b9f6dbd5d", [:mix], [], "hexpm", "30b48089fcc9a648d4bfcd17d51ab799b846c34db0fbf1fd6eb8a5526587510d"},
-  "fss": {:hex, :fss, "0.1.1", "9db2344dbbb5d555ce442ac7c2f82dd975b605b50d169314a20f08ed21e08642", [:mix], [], "hexpm", "78ad5955c7919c3764065b21144913df7515d52e228c09427a004afe9c1a16b0"},
-  "hackney": {:hex, :hackney, "1.25.0", "390e9b83f31e5b325b9f43b76e1a785cbdb69b5b6cd4e079aa67835ded046867", [:rebar3], [{:certifi, "~> 2.15.0", [hex: :certifi, repo: "hexpm", optional: false]}, {:idna, "~> 6.1.0", [hex: :idna, repo: "hexpm", optional: false]}, {:metrics, "~> 1.0.0", [hex: :metrics, repo: "hexpm", optional: false]}, {:mimerl, "~> 1.4", [hex: :mimerl, repo: "hexpm", optional: false]}, {:parse_trans, "3.4.1", [hex: :parse_trans, repo: "hexpm", optional: false]}, {:ssl_verify_fun, "~> 1.1.0", [hex: :ssl_verify_fun, repo: "hexpm", optional: false]}, {:unicode_util_compat, "~> 0.7.1", [hex: :unicode_util_compat, repo: "hexpm", optional: false]}], "hexpm", "7209bfd75fd1f42467211ff8f59ea74d6f2a9e81cbcee95a56711ee79fd6b1d4"},
-  "httpoison": {:hex, :httpoison, "2.2.3", "a599d4b34004cc60678999445da53b5e653630651d4da3d14675fedc9dd34bd6", [:mix], [{:hackney, "~> 1.21", [hex: :hackney, repo: "hexpm", optional: false]}], "hexpm", "fa0f2e3646d3762fdc73edb532104c8619c7636a6997d20af4003da6cfc53e53"},
-  "idna": {:hex, :idna, "6.1.1", "8a63070e9f7d0c62eb9d9fcb360a7de382448200fbbd1b106cc96d3d8099df8d", [:rebar3], [{:unicode_util_compat, "~> 0.7.0", [hex: :unicode_util_compat, repo: "hexpm", optional: false]}], "hexpm", "92376eb7894412ed19ac475e4a86f7b413c1b9fbb5bd16dccd57934157944cea"},
-  "jason": {:hex, :jason, "1.4.4", "b9226785a9aa77b6857ca22832cffa5d5011a667207eb2a0ad56adb5db443b8a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "c5eb0cab91f094599f94d55bc63409236a8ec69a21a67814529e8d5f6cc90b3b"},
-  "kino": {:hex, :kino, "0.17.0", "72f1a2bf691db7b8352bae86b3951fdf9b23619b5d8586cb7cd1e9c2edc8ff9b", [:mix], [{:fss, "~> 0.1.0", [hex: :fss, repo: "hexpm", optional: false]}, {:nx, "~> 0.1", [hex: :nx, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}, {:table, "~> 0.1.2", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "e1ec49a2ebbf622c1675f96b427c565ce02df6725e8f2e8d4a743c8e791bd090"},
+  "jason": {:hex, :jason, "1.4.5", "2e3a008590b0b8d7388c20293e9dcc9cf3e5d642fd2a114e4cbbb52e595d940a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "b0c823996102bcd0239b3c2444eb00409b72f6a140c1950bc8b457d836b30684"},
+  "kino": {:hex, :kino, "0.19.0", "fc8e46fefeb2d083e757633ddd810c499754b7a1f87ba1e92844791c1eca87c2", [:mix], [{:nx, "~> 0.1", [hex: :nx, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}, {:table, "~> 0.1.2", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "195956058730acc397375a398835c7818f92aa01a1d32c03ecb5dffff74c0b8c"},
   "kino_vega_lite": {:hex, :kino_vega_lite, "0.1.13", "03c00405987a2202e4b8014ee55eb7f5727691b3f13d76a3764f6eeccef45322", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: false]}, {:vega_lite, "~> 0.1.8", [hex: :vega_lite, repo: "hexpm", optional: false]}], "hexpm", "00c72bc270e7b9d3c339f726cdab0012fd3f2fc75e36c7548e0f250fe420fa10"},
   "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"},
-  "makeup_elixir": {:hex, :makeup_elixir, "0.16.2", "627e84b8e8bf22e60a2579dad15067c755531fea049ae26ef1020cad58fe9578", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "41193978704763f6bbe6cc2758b84909e62984c7752b3784bd3c218bb341706b"},
-  "makeup_erlang": {:hex, :makeup_erlang, "0.1.5", "e0ff5a7c708dda34311f7522a8758e23bfcd7d8d8068dc312b5eb41c6fd76eba", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "94d2e986428585a21516d7d7149781480013c56e30c6a233534bedf38867a59a"},
-  "metrics": {:hex, :metrics, "1.0.1", "25f094dea2cda98213cecc3aeff09e940299d950904393b2a29d191c346a8486", [:rebar3], [], "hexpm", "69b09adddc4f74a40716ae54d140f93beb0fb8978d8636eaded0c31b6f099f16"},
-  "mimerl": {:hex, :mimerl, "1.4.0", "3882a5ca67fbbe7117ba8947f27643557adec38fa2307490c4c4207624cb213b", [:rebar3], [], "hexpm", "13af15f9f68c65884ecca3a3891d50a7b57d82152792f3e19d88650aa126b144"},
+  "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"},
+  "makeup_erlang": {:hex, :makeup_erlang, "1.0.3", "4252d5d4098da7415c390e847c814bad3764c94a814a0b4245176215615e1035", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "953297c02582a33411ac6208f2c6e55f0e870df7f80da724ed613f10e6706afd"},
   "nimble_csv": {:hex, :nimble_csv, "1.3.0", "b7f998dc62b222bce9596e46f028c7a5af04cb5dde6df2ea197c583227c54971", [:mix], [], "hexpm", "41ccdc18f7c8f8bb06e84164fc51635321e80d5a3b450761c4997d620925d619"},
   "nimble_options": {:hex, :nimble_options, "1.1.1", "e3a492d54d85fc3fd7c5baf411d9d2852922f66e69476317787a7b2bb000a61b", [:mix], [], "hexpm", "821b2470ca9442c4b6984882fe9bb0389371b8ddec4d45a9504f00a66f650b44"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"},
-  "nx": {:hex, :nx, "0.10.0", "128e4a094cb790f663e20e1334b127c1f2a4df54edfb8b13c22757ec33133b4f", [:mix], [{:complex, "~> 0.6", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "3db8892c124aeee091df0e6fbf8e5bf1b81f502eb0d4f5ba63e6378ebcae7da4"},
-  "parse_trans": {:hex, :parse_trans, "3.4.1", "6e6aa8167cb44cc8f39441d05193be6e6f4e7c2946cb2759f015f8c56b76e5ff", [:rebar3], [], "hexpm", "620a406ce75dada827b82e453c19cf06776be266f5a67cff34e1ef2cbb60e49a"},
-  "rustler_precompiled": {:hex, :rustler_precompiled, "0.8.3", "4e741024b0b097fe783add06e53ae9a6f23ddc78df1010f215df0c02915ef5a8", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "c23f5f33cb6608542de4d04faf0f0291458c352a4648e4d28d17ee1098cddcc4"},
+  "nx": {:hex, :nx, "0.12.1", "6e9fee43a77646d04faad2ba4e449b9e270c6b23e413f203cb4d81e71c0a617f", [:mix], [{:complex, "~> 0.7", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "ee80f6ae898f68bbfe7f30216b06aab10096231ec99ded1208a827256b23d0fb"},
+  "rustler_precompiled": {:hex, :rustler_precompiled, "0.8.4", "700a878312acfac79fb6c572bb8b57f5aae05fe1cf70d34b5974850bbf2c05bf", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "3b33d99b540b15f142ba47944f7a163a25069f6d608783c321029bc1ffb09514"},
   "scidata": {:hex, :scidata, "0.1.11", "fe3358bac7d740374b4f2a7eff6a1cb02e5ee7f87f7cdb1e8648ad93c533165f", [:mix], [{:castore, "~> 0.1", [hex: :castore, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.1", [hex: :nimble_csv, repo: "hexpm", optional: false]}, {:stb_image, "~> 0.4", [hex: :stb_image, repo: "hexpm", optional: true]}], "hexpm", "90873337a9d5fe880d640517efa93d3c07e46c8ba436de44117f581800549f93"},
-  "ssl_verify_fun": {:hex, :ssl_verify_fun, "1.1.7", "354c321cf377240c7b8716899e182ce4890c5938111a1296add3ec74cf1715df", [:make, :mix, :rebar3], [], "hexpm", "fe4c190e8f37401d30167c8c405eda19469f34577987c76dde613e838bbc67f8"},
   "table": {:hex, :table, "0.1.2", "87ad1125f5b70c5dea0307aa633194083eb5182ec537efc94e96af08937e14a8", [:mix], [], "hexpm", "7e99bc7efef806315c7e65640724bf165c3061cdc5d854060f74468367065029"},
   "telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"},
-  "unicode_util_compat": {:hex, :unicode_util_compat, "0.7.1", "a48703a25c170eedadca83b11e88985af08d35f37c6f664d6dcfb106a97782fc", [:rebar3], [], "hexpm", "b3a917854ce3ae233619744ad1e0102e05673136776fb2fa76234f3e03b23642"},
   "vega_lite": {:hex, :vega_lite, "0.1.11", "2b261d21618f6fa9f63bb4542f0262982d2e40aea3f83e935788fe172902b3c2", [:mix], [{:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "d18c3f11369c14bdf36ab53010c06bf5505c221cbcb32faac7420cf6926b3c50"},
   "vega_lite_convert": {:hex, :vega_lite_convert, "1.0.1", "1cc0309998c10bce9d944ae631938a433e9cad4ccf7344f9a192d9ddcab1bd93", [:mix], [{:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: false]}, {:rustler, ">= 0.0.0", [hex: :rustler, repo: "hexpm", optional: true]}, {:rustler_precompiled, "~> 0.4", [hex: :rustler_precompiled, repo: "hexpm", optional: false]}, {:vega_lite, ">= 0.0.0", [hex: :vega_lite, repo: "hexpm", optional: false]}], "hexpm", "1e12b4ef3943510d3597e4876a10c6ab333b8ee71204b5e571605d80aa3cbd76"},
 }
diff --git a/scripts/check_xgboost_c_api.sh b/scripts/check_xgboost_c_api.sh
new file mode 100755
index 0000000..6ed3fa5
--- /dev/null
+++ b/scripts/check_xgboost_c_api.sh
@@ -0,0 +1,186 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+extract_signatures() {
+  local header="$1"
+
+  awk '
+    /XG[A-Za-z0-9_]+[[:space:]]*\(/ {
+      in_decl=1
+      decl=$0
+      if ($0 ~ /;/) {
+        print decl
+        in_decl=0
+        decl=""
+      }
+      next
+    }
+    in_decl {
+      decl=decl " " $0
+      if ($0 ~ /;/) {
+        print decl
+        in_decl=0
+        decl=""
+      }
+    }
+  ' "$header" \
+    | sed -E 's/[[:space:]]+/ /g; s/^[[:space:]]+//; s/[[:space:]]+$//' \
+    | awk '
+        {
+          if (match($0, /XG[A-Za-z0-9_]+[[:space:]]*\(/)) {
+            name = substr($0, RSTART, RLENGTH)
+            gsub(/[[:space:]]*\($/, "", name)
+            print name "\t" $0
+          }
+        }
+      '
+}
+
+compare_mode() {
+  if [[ $# -ne 2 ]]; then
+    echo "Usage: $0 --compare <old-xgboost-include-dir> <new-xgboost-include-dir>"
+    exit 2
+  fi
+
+  local old_header="$1/xgboost/c_api.h"
+  local new_header="$2/xgboost/c_api.h"
+
+  if [[ ! -f "$old_header" ]]; then
+    echo "XGBoost C API header not found: $old_header"
+    exit 2
+  fi
+
+  if [[ ! -f "$new_header" ]]; then
+    echo "XGBoost C API header not found: $new_header"
+    exit 2
+  fi
+
+  old_tmp=$(mktemp)
+  new_tmp=$(mktemp)
+  trap 'rm -f "${old_tmp:-}" "${new_tmp:-}"' EXIT
+
+  extract_signatures "$old_header" | sort -u > "$old_tmp"
+  extract_signatures "$new_header" | sort -u > "$new_tmp"
+
+  declare -A old_map=()
+  declare -A new_map=()
+
+  while IFS=$'\t' read -r name sig; do
+    [[ -n "${name:-}" ]] && old_map["$name"]="$sig"
+  done < "$old_tmp"
+
+  while IFS=$'\t' read -r name sig; do
+    [[ -n "${name:-}" ]] && new_map["$name"]="$sig"
+  done < "$new_tmp"
+
+  removed=()
+  added=()
+  changed=()
+
+  for name in "${!old_map[@]}"; do
+    if [[ -z "${new_map[$name]+x}" ]]; then
+      removed+=("$name")
+    elif [[ "${old_map[$name]}" != "${new_map[$name]}" ]]; then
+      changed+=("$name")
+    fi
+  done
+
+  for name in "${!new_map[@]}"; do
+    if [[ -z "${old_map[$name]+x}" ]]; then
+      added+=("$name")
+    fi
+  done
+
+  if (( ${#removed[@]} == 0 && ${#changed[@]} == 0 && ${#added[@]} == 0 )); then
+    echo "No C API declaration differences found between versions."
+    return 0
+  fi
+
+  if (( ${#removed[@]} > 0 )); then
+    printf 'Removed symbols (%d):\n' "${#removed[@]}"
+    printf '  - %s\n' "${removed[@]}"
+  fi
+
+  if (( ${#changed[@]} > 0 )); then
+    printf 'Changed signatures (%d):\n' "${#changed[@]}"
+    for name in "${changed[@]}"; do
+      printf '  - %s\n' "$name"
+      printf '      old: %s\n' "${old_map[$name]}"
+      printf '      new: %s\n' "${new_map[$name]}"
+    done
+  fi
+
+  if (( ${#added[@]} > 0 )); then
+    printf 'Added symbols (%d):\n' "${#added[@]}"
+    printf '  - %s\n' "${added[@]}"
+  fi
+
+  if (( ${#removed[@]} > 0 || ${#changed[@]} > 0 )); then
+    return 1
+  fi
+
+  return 0
+}
+
+if [[ "${1:-}" == "--compare" ]]; then
+  shift
+  compare_mode "$@"
+  exit $?
+fi
+
+if [[ $# -lt 1 ]]; then
+  echo "Usage: $0 <xgboost-install-include-dir> [xgboost-shared-lib-path]"
+  echo "       $0 --compare <old-xgboost-include-dir> <new-xgboost-include-dir>"
+  exit 2
+fi
+
+include_dir="$1"
+shared_lib="${2:-}"
+header="${include_dir}/xgboost/c_api.h"
+
+if [[ ! -f "$header" ]]; then
+  echo "XGBoost C API header not found: $header"
+  exit 2
+fi
+
+mapfile -t used_symbols < <(
+  grep -RhoE '\bXG[A-Za-z0-9_]+[[:space:]]*\(' c/exgboost/src/*.c \
+    | sed -E 's/[[:space:]]*\($//' \
+    | sort -u
+)
+
+missing_in_header=()
+for symbol in "${used_symbols[@]}"; do
+  if ! grep -qE "\\b${symbol}[[:space:]]*\\(" "$header"; then
+    missing_in_header+=("$symbol")
+  fi
+done
+
+if (( ${#missing_in_header[@]} > 0 )); then
+  echo "Symbols used by exgboost NIF but missing from XGBoost C API header:"
+  printf '  - %s\n' "${missing_in_header[@]}"
+  exit 1
+fi
+
+if [[ -n "$shared_lib" && -f "$shared_lib" && "$(command -v nm || true)" != "" ]]; then
+  if [[ "$(uname -s)" == "Darwin" ]]; then
+    mapfile -t exported_symbols < <(nm -gU "$shared_lib" | awk '{print $3}' | sed 's/^_//' | sort -u)
+  else
+    mapfile -t exported_symbols < <(nm -D --defined-only "$shared_lib" | awk '{print $3}' | sort -u)
+  fi
+
+  missing_in_lib=()
+  for symbol in "${used_symbols[@]}"; do
+    if ! printf '%s\n' "${exported_symbols[@]}" | grep -qx "$symbol"; then
+      missing_in_lib+=("$symbol")
+    fi
+  done
+
+  if (( ${#missing_in_lib[@]} > 0 )); then
+    echo "Symbols declared in header but not exported by shared library:"
+    printf '  - %s\n' "${missing_in_lib[@]}"
+    exit 1
+  fi
+fi
+
+echo "XGBoost C API compatibility check passed (${#used_symbols[@]} symbols)."
diff --git a/test/exgboost_test.exs b/test/exgboost_test.exs
index fcdfec0..9c3357a 100644
--- a/test/exgboost_test.exs
+++ b/test/exgboost_test.exs
@@ -80,7 +80,7 @@ defmodule EXGBoostTest do
     assert Booster.get_boosted_rounds(booster) == num_boost_round
   end
 
-  test "predict", context do
+  test "predict matches inplace_predict", context do
     nrows = :rand.uniform(10)
     ncols = :rand.uniform(10)
     {x, new_key} = Nx.Random.normal(context.key, 0, 1, shape: {nrows, ncols})
@@ -95,6 +95,59 @@ defmodule EXGBoostTest do
     assert inplace_preds_no_proxy.shape == y.shape
   end
 
+  test "predict is stable across repeated calls", context do
+    nrows = 12
+    ncols = 5
+    {x, new_key} = Nx.Random.normal(context.key, 0, 1, shape: {nrows, ncols})
+    {y, _new_key} = Nx.Random.normal(new_key, 0, 1, shape: {nrows})
+
+    booster =
+      EXGBoost.train(x, y,
+        num_boost_rounds: 25,
+        tree_method: :hist,
+        eval_metric: :rmse
+      )
+
+    first_preds = EXGBoost.predict(booster, x)
+    second_preds = EXGBoost.predict(booster, x)
+    third_preds = EXGBoost.predict(booster, x)
+
+    assert Nx.all_close(first_preds, second_preds)
+    assert Nx.all_close(first_preds, third_preds)
+    assert Nx.all_close(second_preds, third_preds)
+  end
+
+  test "predict first call matches subsequent confidence values", context do
+    nrows = 30
+    ncols = 5
+    {x, new_key} = Nx.Random.normal(context.key, 0, 1, shape: {nrows, ncols})
+    {y, _new_key} = Nx.Random.normal(new_key, 0, 1, shape: {nrows})
+
+    booster =
+      EXGBoost.train(x, y,
+        num_boost_rounds: 30,
+        tree_method: :hist,
+        eval_metric: :rmse
+      )
+
+    sample = Nx.slice_along_axis(x, 0, 1, axis: 0)
+
+    [baseline_confidence] =
+      EXGBoost.predict(booster, sample)
+      |> Nx.to_flat_list()
+
+    for i <- 1..10 do
+      [confidence] =
+        EXGBoost.predict(booster, sample)
+        |> Nx.to_flat_list()
+
+      assert_in_delta confidence,
+                      baseline_confidence,
+                      1.0e-9,
+                      "Expected confidence to match baseline #{baseline_confidence}, got #{confidence} for iteration #{i}"
+    end
+  end
+
   test "predict with container", context do
     nrows = :rand.uniform(10)
     ncols = :rand.uniform(10)
@@ -166,9 +219,6 @@ defmodule EXGBoostTest do
                      )
                    end)
                  end
-
-    refute is_nil(booster.best_iteration)
-    refute is_nil(booster.best_score)
   end
 
   test "eval with multiple metrics", context do
@@ -280,6 +330,24 @@ defmodule EXGBoostTest do
     assert is_struct(bst, EXGBoost.Booster)
   end
 
+  test "load_model accepts model artifacts produced by dump_weights", context do
+    nrows = :rand.uniform(10)
+    ncols = :rand.uniform(10)
+    {x, _new_key} = Nx.Random.normal(context[:key], 0, 1, shape: {nrows, ncols})
+    {y, _new_key} = Nx.Random.normal(context[:key], 0, 1, shape: {nrows})
+
+    booster =
+      EXGBoost.train(x, y,
+        num_boost_rounds: 10,
+        tree_method: :hist,
+        eval_metric: :rmse
+      )
+
+    buffer = EXGBoost.dump_weights(booster)
+    bst = EXGBoost.load_model(buffer)
+    assert is_struct(bst, EXGBoost.Booster)
+  end
+
   test "serialize and deserialize weights to and from buffer", context do
     nrows = :rand.uniform(10)
     ncols = :rand.uniform(10)
@@ -329,6 +397,37 @@ defmodule EXGBoostTest do
     assert EXGBoost.ArrayInterface.get_tensor(array_interface) == tensor
   end
 
+  test "array interface from_map ignores optional keys" do
+    arr_int =
+      EXGBoost.ArrayInterface.from_map(%{
+        "typestr" => "<f4",
+        "shape" => [2, 2],
+        "data" => [123, true],
+        "version" => 3,
+        "strides" => nil,
+        "descr" => [["", "<f4"]]
+      })
+
+    assert arr_int.typestr == "<f4"
+    assert arr_int.shape == {2, 2}
+    assert arr_int.address == 123
+    assert arr_int.readonly == true
+    assert arr_int.version == 3
+  end
+
+  test "array interface get_tensor raises on unsupported endianness" do
+    assert_raise ArgumentError, ~r/Unsupported endianness/, fn ->
+      EXGBoost.ArrayInterface.get_tensor(%EXGBoost.ArrayInterface{
+        typestr: ">f4",
+        shape: {1},
+        address: 1,
+        readonly: true,
+        tensor: nil,
+        binary: <<>>
+      })
+    end
+  end
+
   describe "errors" do
     setup %{key: key0} do
       {nrows, ncols} = {10, 10}
diff --git a/test/nif_contract_test.exs b/test/nif_contract_test.exs
new file mode 100644
index 0000000..dfc0279
--- /dev/null
+++ b/test/nif_contract_test.exs
@@ -0,0 +1,23 @@
+defmodule NifContractTest do
+  use ExUnit.Case, async: true
+
+  @nif_c_file Path.expand("../c/exgboost/src/exgboost.c", __DIR__)
+
+  test "Elixir NIF stubs match C registered NIF functions" do
+    c_source = File.read!(@nif_c_file)
+
+    c_exports =
+      Regex.scan(~r/\{"([a-z0-9_]+)",\s*(\d+),/i, c_source)
+      |> Enum.map(fn [_, name, arity] -> {String.to_atom(name), String.to_integer(arity)} end)
+      |> MapSet.new()
+
+    elixir_exports =
+      EXGBoost.NIF.__info__(:functions)
+      |> Enum.reject(fn {name, _arity} -> name in [:module_info, :__info__, :on_load] end)
+      |> MapSet.new()
+
+    assert c_exports == elixir_exports,
+           "NIF API mismatch. Missing in C: #{inspect(MapSet.difference(elixir_exports, c_exports))}. " <>
+             "Missing in Elixir: #{inspect(MapSet.difference(c_exports, elixir_exports))}"
+  end
+end
diff --git a/test/nif_test.exs b/test/nif_test.exs
index 383418d..39d75c3 100644
--- a/test/nif_test.exs
+++ b/test/nif_test.exs
@@ -7,6 +7,12 @@ defmodule NifTest do
     assert EXGBoost.NIF.xgboost_version() |> unwrap!() != :error
   end
 
+  test "unwrap! raises string error for charlist reasons" do
+    assert_raise ArgumentError, "boom", fn ->
+      unwrap!({:error, ~c"boom"})
+    end
+  end
+
   test "build_info" do
     assert EXGBoost.NIF.xgboost_build_info() |> unwrap!() != :error
   end