diff --git a/synthimpute/rf_impute.py b/synthimpute/rf_impute.py
index bdec8fc..06b22a4 100644
--- a/synthimpute/rf_impute.py
+++ b/synthimpute/rf_impute.py
@@ -45,7 +45,7 @@ def rf_quantile(m, X, q):
 
 
 def rf_impute(x_train, y_train, x_new, x_cols=None, random_state=None,
-              **kwargs):
+              sample_weight_train=None, **kwargs):
     """Impute labels from a training set to a new data set using 
        random forests quantile regression.
        
@@ -57,6 +57,8 @@ def rf_impute(x_train, y_train, x_new, x_cols=None, random_state=None,
             x_train (these must also be in x_new).
         random_state: Optional random seed passed to RandomForestRegressor and
             for uniform distribution of quantiles.
+        sample_weight_train: Vector indicating the weights associated with each
+            row of x_train/y_train. Defaults to None.
         **kwargs: Other args passed to RandomForestRegressor, e.g. 
             `n_estimators=50`.  rf_impute uses all RandomForestRegressor
             defaults unless otherwise specified.
@@ -65,7 +67,10 @@ def rf_impute(x_train, y_train, x_new, x_cols=None, random_state=None,
         Imputed labels for new_x.
     """
     rf = ensemble.RandomForestRegressor(random_state=random_state, **kwargs)
-    rf.fit(x_train, y_train)
+    if sample_weight_train is None:
+        rf.fit(x_train, y_train)
+    else:
+        rf.fit(x_train, y_train, sample_weight=sample_weight_train)
     if random_state is not None:
         np.random.seed(random_state)
     quantiles = np.random.rand(y_train.size)  # Uniform distribution.

Add observation weights to rf_impute
`sample_weight` in `fit` should do it: https://scikit-learn.org/stable/modules/generated/sklearn.ensemble.RandomForestRegressor.html#sklearn.ensemble.RandomForestRegressor.fit
diff --git a/synthimpute/tests/test_rf_impute.py b/synthimpute/tests/test_rf_impute.py
index c02b6f0..ce99542 100644
--- a/synthimpute/tests/test_rf_impute.py
+++ b/synthimpute/tests/test_rf_impute.py
@@ -12,7 +12,14 @@ def test_rf_impute():
                       'x2': np.random.randn(n)})
     # Construct example relationship.
     y = x.x1 + np.power(x.x2, 3) + np.random.randn(n)
-    si.rf_impute(x.iloc[:N_TRAIN], y.iloc[:N_TRAIN], x.iloc[N_TRAIN:])
+    # Split into test and train.
+    x_train = x.iloc[:N_TRAIN]
+    y_train = y.iloc[:N_TRAIN]
+    x_test = x.iloc[N_TRAIN:]
+    si.rf_impute(x_train, y_train, x_test)
     # Try with some args.
-    si.rf_impute(x.iloc[:N_TRAIN], y.iloc[:N_TRAIN], x.iloc[N_TRAIN:],
+    si.rf_impute(x_train, y_train, x_test,
                  random_state=10, n_estimators=200)
+    # Try with sample_weight_train.
+    si.rf_impute(x_train, y_train, x_test,
+                 sample_weight_train=np.random.randn(N_TRAIN))

pip install -e .
parse_log_pytest
numpy>=1.16.0 pandas>=1.0.0 scikit-learn>=0.22.0 statsmodels>=0.10.0
pytest pytest-cov pytest-xdist pytest-mock pytest-asyncio
apt-get update apt-get install -y gcc
3.9
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
Bottleneck @ file:///croot/bottleneck_1731058641041/work
coverage==7.8.0
exceptiongroup==1.2.2
execnet==2.1.1
iniconfig==2.1.0
joblib @ file:///croot/joblib_1718217211762/work
numexpr @ file:///croot/numexpr_1730215937391/work
numpy @ file:///croot/numpy_and_numpy_base_1736283260865/work/dist/numpy-2.0.2-cp39-cp39-linux_x86_64.whl#sha256=3387e3e62932fa288bc18e8f445ce19e998b418a65ed2064dd40a054f976a6c7
packaging @ file:///croot/packaging_1734472117206/work
pandas @ file:///croot/pandas_1732735089971/work/dist/pandas-2.2.3-cp39-cp39-linux_x86_64.whl#sha256=0a51ed2e81ab863e3d00ed6c5049192ce578ecb38fb467d2f9a6585d3c25f666
patsy @ file:///croot/patsy_1738159930729/work
pluggy==1.5.0
pytest==8.3.5
pytest-asyncio==0.26.0
pytest-cov==6.0.0
pytest-mock==3.14.0
pytest-xdist==3.6.1
python-dateutil @ file:///croot/python-dateutil_1716495738603/work
pytz @ file:///croot/pytz_1713974312559/work
scikit-learn @ file:///croot/scikit-learn_1737988764076/work
scipy @ file:///croot/scipy_1733756309941/work/dist/scipy-1.13.1-cp39-cp39-linux_x86_64.whl#sha256=3b247b926209f2d9f719ebae39faf3ff891b2596150ed8f8349adfc3eb19441c
six @ file:///tmp/build/80754af9/six_1644875935023/work
statsmodels @ file:///croot/statsmodels_1734018870787/work
-e git+https://github.com/MaxGhenis/synthimpute.git@f668c4ab967e2a42bc0bed668d59efc681bf2485#egg=synthimpute
threadpoolctl @ file:///croot/threadpoolctl_1719407800858/work
tomli==2.2.1
typing_extensions==4.13.0
tzdata @ file:///croot/python-tzdata_1690578112552/work

name: synthimpute
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - blas=1.0=openblas
  - bottleneck=1.4.2=py39ha9d4c09_0
  - ca-certificates=2025.2.25=h06a4308_0
  - joblib=1.4.2=py39h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgfortran-ng=11.2.0=h00389a5_1
  - libgfortran5=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libopenblas=0.3.21=h043d6bf_0
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - numexpr=2.10.1=py39hd28fd6d_0
  - numpy=2.0.2=py39heeff2f4_0
  - numpy-base=2.0.2=py39h8a23956_0
  - openssl=3.0.16=h5eee18b_0
  - packaging=24.2=py39h06a4308_0
  - pandas=2.2.3=py39h6a678d5_0
  - patsy=1.0.1=py39h06a4308_0
  - pip=25.0=py39h06a4308_0
  - pybind11-abi=4=hd3eb1b0_1
  - python=3.9.21=he870216_1
  - python-dateutil=2.9.0post0=py39h06a4308_2
  - python-tzdata=2023.3=pyhd3eb1b0_0
  - pytz=2024.1=py39h06a4308_0
  - readline=8.2=h5eee18b_0
  - scikit-learn=1.6.1=py39h6a678d5_0
  - scipy=1.13.1=py39heeff2f4_1
  - setuptools=72.1.0=py39h06a4308_0
  - six=1.16.0=pyhd3eb1b0_1
  - sqlite=3.45.3=h5eee18b_0
  - statsmodels=0.14.4=py39h5eee18b_0
  - threadpoolctl=3.5.0=py39h2f386ee_0
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - coverage==7.8.0
      - exceptiongroup==1.2.2
      - execnet==2.1.1
      - iniconfig==2.1.0
      - pluggy==1.5.0
      - pytest==8.3.5
      - pytest-asyncio==0.26.0
      - pytest-cov==6.0.0
      - pytest-mock==3.14.0
      - pytest-xdist==3.6.1
      - tomli==2.2.1
      - typing-extensions==4.13.0
prefix: /opt/conda/envs/synthimpute

synthimpute/tests/test_rf_impute.py::test_rf_impute
synthimpute/rf_impute.py
synthimpute/rf_impute.py