Skip to content

Commit

Permalink
Make rust Tensor archetype eager serialized (#8801)
Browse files Browse the repository at this point in the history
### Related

* sister PR to..
	* #8789
	* #8785
	* #8793
* missed piece of #7245

### What

Ports the Tensor archetype in rust to the new eager serialized
interface.
Unfortunately this meant I had to remove some direct access methods of
the underlying tensor data. Curiously, this didn't affect any of our
test/snippet/example code.
While doing so I also fixed some wording issues in the (very similar)
C++ implementation of `with_dim_names`
  • Loading branch information
Wumpf authored Jan 24, 2025
1 parent 7436dbf commit 061ce27
Show file tree
Hide file tree
Showing 8 changed files with 177 additions and 133 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -6562,6 +6562,7 @@ dependencies = [
"rayon",
"re_build_tools",
"re_byte_size",
"re_error",
"re_format",
"re_log",
"re_log_types",
Expand Down
3 changes: 2 additions & 1 deletion crates/store/re_types/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,10 @@ testing = []
[dependencies]
# Rerun
re_byte_size.workspace = true
re_error.workspace = true
re_format.workspace = true
re_log.workspace = true
re_log_types.workspace = true
re_log.workspace = true
re_tracing.workspace = true
re_types_core.workspace = true
re_video = { workspace = true, optional = true }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ namespace rerun.archetypes;
///
/// \example archetypes/tensor_simple title="Simple tensor" image="https://static.rerun.io/tensor_simple/baacb07712f7b706e3c80e696f70616c6c20b367/1200w.png"
table Tensor (
// TODO(#7245): "attr.rust.archetype_eager",
"attr.rust.archetype_eager",
"attr.rust.derive": "PartialEq",
"attr.docs.category": "Image & tensor",
"attr.docs.view_types": "TensorView, BarChartView: for 1D tensors"
Expand Down
169 changes: 118 additions & 51 deletions crates/store/re_types/src/archetypes/tensor.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

111 changes: 37 additions & 74 deletions crates/store/re_types/src/archetypes/tensor_ext.rs
Original file line number Diff line number Diff line change
@@ -1,25 +1,19 @@
use crate::datatypes::TensorData;
use crate::{
components::{self},
datatypes::TensorData,
};

use re_types_core::ArrowString;
use re_types_core::{try_serialize_field, ArrowString, Loggable};

use super::Tensor;

impl Tensor {
/// Accessor to the underlying [`TensorData`].
pub fn data(&self) -> &TensorData {
&self.data.0
}

/// Try to construct a [`Tensor`] from anything that can be converted into [`TensorData`]
///
/// This is useful for constructing a tensor from an ndarray.
pub fn try_from<T: TryInto<TensorData>>(data: T) -> Result<Self, T::Error> {
let data: TensorData = data.try_into()?;

Ok(Self {
data: data.into(),
value_range: None,
})
Ok(Self::new(data))
}

/// Update the `names` of the contained [`TensorData`] dimensions.
Expand All @@ -32,7 +26,35 @@ impl Tensor {
mut self,
names: impl IntoIterator<Item = impl Into<ArrowString>>,
) -> Self {
self.data.0 = self.data.0.with_dim_names(names);
if let Some(data) = self.data.take() {
match components::TensorData::from_arrow(&data.array) {
Ok(tensor_data) => {
if tensor_data.len() > 1 {
re_log::warn!(
"Can't set dimension names on a tensor archetype with multiple tensor data instances."
);
return self;
}
let Some(tensor_data) = tensor_data.into_iter().next() else {
re_log::warn!(
"Can't set dimension names on a tensor archetype without any tensor data instances."
);
return self;
};

self.data = try_serialize_field::<components::TensorData>(
Self::descriptor_data(),
[components::TensorData(tensor_data.0.with_dim_names(names))],
);
}
Err(err) => re_log::warn!(
"Failed to read arrow tensor data: {}",
re_error::format_ref(&err)
),
}
} else {
re_log::warn!("Can't set names on a tensor that doesn't have any data");
}
self
}
}
Expand All @@ -45,10 +67,7 @@ impl Tensor {
pub fn from_image(
image: impl Into<image::DynamicImage>,
) -> Result<Self, crate::tensor_data::TensorImageLoadError> {
TensorData::from_image(image).map(|data| Self {
data: data.into(),
value_range: None,
})
TensorData::from_image(image).map(Self::new)
}

/// Construct a tensor from [`image::DynamicImage`].
Expand All @@ -57,62 +76,6 @@ impl Tensor {
pub fn from_dynamic_image(
image: image::DynamicImage,
) -> Result<Self, crate::tensor_data::TensorImageLoadError> {
TensorData::from_dynamic_image(image).map(|data| Self {
data: data.into(),
value_range: None,
})
}
}

impl AsRef<TensorData> for Tensor {
#[inline(always)]
fn as_ref(&self) -> &TensorData {
&self.data
}
}

impl std::ops::Deref for Tensor {
type Target = TensorData;

#[inline(always)]
fn deref(&self) -> &TensorData {
&self.data
TensorData::from_dynamic_image(image).map(Self::new)
}
}

impl std::borrow::Borrow<TensorData> for Tensor {
#[inline(always)]
fn borrow(&self) -> &TensorData {
&self.data
}
}

// ----------------------------------------------------------------------------
// Make it possible to create an ArrayView directly from a Tensor.

macro_rules! forward_array_views {
($type:ty, $alias:ty) => {
impl<'a> TryFrom<&'a $alias> for ::ndarray::ArrayViewD<'a, $type> {
type Error = crate::tensor_data::TensorCastError;

#[inline]
fn try_from(value: &'a $alias) -> Result<Self, Self::Error> {
value.data().try_into()
}
}
};
}

forward_array_views!(u8, Tensor);
forward_array_views!(u16, Tensor);
forward_array_views!(u32, Tensor);
forward_array_views!(u64, Tensor);

forward_array_views!(i8, Tensor);
forward_array_views!(i16, Tensor);
forward_array_views!(i32, Tensor);
forward_array_views!(i64, Tensor);

forward_array_views!(half::f16, Tensor);
forward_array_views!(f32, Tensor);
forward_array_views!(f64, Tensor);
Loading

0 comments on commit 061ce27

Please sign in to comment.