Use generic names in multidimensional add_features test

This commit is contained in:
felixmin
2026-03-21 12:43:12 +01:00
committed by CarolinePascal
parent e0e03887e0
commit 910aac4c03
+8 -8
View File
@@ -381,7 +381,7 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
def test_add_features_with_multidimensional_values(sample_dataset, tmp_path): def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
"""Test adding a multi-dimensional per-frame feature.""" """Test adding a multi-dimensional per-frame feature."""
num_frames = sample_dataset.meta.total_frames num_frames = sample_dataset.meta.total_frames
latent_values = np.random.randn(num_frames, 1, 4).astype(np.float32) feature_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
feature_info = { feature_info = {
"dtype": "float32", "dtype": "float32",
@@ -389,7 +389,7 @@ def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
"names": None, "names": None,
} }
features = { features = {
"latent_vectors": (latent_values, feature_info), "multi_dim_feature": (feature_values, feature_info),
} }
with ( with (
@@ -397,19 +397,19 @@ def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download, patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
): ):
mock_get_safe_version.return_value = "v3.0" mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "with_latents") mock_snapshot_download.return_value = str(tmp_path / "with_multi_dim_feature")
new_dataset = add_features( new_dataset = add_features(
dataset=sample_dataset, dataset=sample_dataset,
features=features, features=features,
output_dir=tmp_path / "with_latents", output_dir=tmp_path / "with_multi_dim_feature",
) )
assert "latent_vectors" in new_dataset.meta.features assert "multi_dim_feature" in new_dataset.meta.features
sample_item = new_dataset[0] sample_item = new_dataset[0]
assert "latent_vectors" in sample_item assert "multi_dim_feature" in sample_item
assert isinstance(sample_item["latent_vectors"], torch.Tensor) assert isinstance(sample_item["multi_dim_feature"], torch.Tensor)
assert tuple(sample_item["latent_vectors"].shape) == (1, 4) assert tuple(sample_item["multi_dim_feature"].shape) == (1, 4)
def test_add_existing_feature(sample_dataset, tmp_path): def test_add_existing_feature(sample_dataset, tmp_path):