diff --git a/README.md b/README.md index e0f9464..cffb6d2 100644 --- a/README.md +++ b/README.md @@ -8,8 +8,8 @@ [![Discord](https://img.shields.io/discord/1475229838754316502?label=discord)](https://discord.gg/vBgXHWCeah) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zspacelabs/bunsen) -`bunsen` aims to be a "batteries included" complementary -community standard library for extending the [burn](https://burn.dev) tensor library. +`bunsen` aims to be a "batteries included" complementary community standard library for extending +the [burn](https://burn.dev) tensor library. # Book @@ -19,28 +19,24 @@ Read the [bunsen book](https://zspacelabs.ai/bunsen/book) ## Public / API Crates -* [`bunsen-firehose`](crates/bunsen-firehose) — a columnar dataloader / - processing pipeline, with a burn batcher bridge. +* [`bunsen-firehose`](crates/bunsen-firehose) — a columnar dataloader / processing pipeline, with a burn batcher bridge. ## Utility Crates * [`bunsen-contracts-macros`](crates/bunsen-contracts-macros) — the - `shape_contract![]` proc-macro backing `bunsen`'s runtime tensor-shape - contracts. + `shape_contract![]` proc-macro backing `bunsen`'s runtime tensor-shape contracts. ## Experimental Crates -These represent complex-interface + work-in-progress, unstable interface -extensions to `bunsen`; particulary those which incur large dependencies -or are not yet ready for general consumption. +These represent complex-interface + work-in-progress, unstable interface extensions to `bunsen`; particulary those which +incur large dependencies or are not yet ready for general consumption. -* [`bunsen-firehose-image`](crates/bunsen-firehose-image) — image loading, - augmentation, and tensor-conversion operators for `bunsen-firehose`. -* [`bunsen`](crates/bunsen) — the main "batteries included" library extending - burn: model blocks, kits, ops, contracts, and support tooling. -* [`bunsen-preview-chat-dataloader`](crates/bunsen-preview-chat-dataloader) — - *(preview)* an Arrow-backed chat dataloader with tokenization for LLM - training. +* [`bunsen-firehose-image`](crates/bunsen-firehose-image) — image loading, augmentation, and tensor-conversion operators + for `bunsen-firehose`. +* [`bunsen`](crates/bunsen) — the main "batteries included" library extending burn: model blocks, kits, ops, contracts, + and support tooling. +* [`bunsen-preview-chat-dataloader`](crates/bunsen-preview-chat-dataloader) — *(preview)* an Arrow-backed chat + dataloader with tokenization for LLM training. # API Examples @@ -50,11 +46,10 @@ A "good parts" survey of some of `bunsen`'s features. See the ## Shape Contracts -`bunsen::contracts` provides allocation-free, always-on runtime tensor-shape -contracts. A contract pairs paper-style shape notation with runtime checking: -it asserts that a tensor's shape matches a declared pattern *and* unpacks named -dimensions for downstream use, catching shape errors at their source. Single -checks run in ~160 ns; the amortized periodic variants average a few ns. +`bunsen::contracts` provides allocation-free, always-on runtime tensor-shape contracts. A contract pairs paper-style +shape notation with runtime checking: +it asserts that a tensor's shape matches a declared pattern *and* unpacks named dimensions for downstream use, catching +shape errors at their source. Single checks run in ~160 ns; the amortized periodic variants average a few ns. ```rust use bunsen::contracts::*; @@ -79,8 +74,8 @@ let [b, h_wins, w_wins, c] = unpack_shape_contract!( ); ``` -In hot loops, use `assert_shape_contract_periodically!` to amortize the check -via exponential backoff while still catching regressions: +In hot loops, use `assert_shape_contract_periodically!` to amortize the check via exponential backoff while still +catching regressions: ```rust use bunsen::contracts::*; @@ -94,10 +89,9 @@ assert_shape_contract_periodically!( ## TensorData Index Views -`TensorDataIndexView` and `TensorDataIndexMutView` wrap burn's low-level -`TensorData` to give ergonomic multi-dimensional element access via bracket -notation — `view[&[i, j]]` — instead of manually flattening indices. The views -deref to the underlying `TensorData`, so `.shape` and friends are right there. +`TensorDataView` and `TensorDataViewMut` wrap burn's low-level +`TensorData` to give ergonomic multi-dimensional element access via bracket notation — `view[&[i, j]]` — instead of +manually flattening indices. The views deref to the underlying `TensorData`, so `.shape` and friends are right there. Handy for inspecting or patching raw tensor data without building full tensors. ```rust @@ -130,10 +124,9 @@ assert_eq!(view[&[0, 0]], 10.0); ## XML Module Reflection `bunsen::burner::module::reflection::XmlModuleTree` turns any burn `Module` -into a queryable XML meta-description of its structure. This enables -type-erased introspection and XPath-style parameter selection — e.g. "every -rank-2 weight under the transformer blocks" — which is exactly what you need to -slice a model into parameter groups for per-group optimizers. +into a queryable XML meta-description of its structure. This enables type-erased introspection and XPath-style parameter +selection — e.g. "every rank-2 weight under the transformer blocks" — which is exactly what you need to slice a model +into parameter groups for per-group optimizers. Take a small container module: @@ -172,8 +165,8 @@ let matrix_params = mtree .to_param_ids() ?; ``` -The dumped structure mirrors the module's fields, with each `@name` taken from -the struct field and a stable `param_id` per tensor: +The dumped structure mirrors the module's fields, with each `@name` taken from the struct field and a stable `param_id` +per tensor: ```xml @@ -199,8 +192,8 @@ the struct field and a stable `param_id` per tensor: ## Blocks & Ops -`bunsen::blocks` is a library of `burn::module::Module` building blocks (stateful -layers with trainable parameters), and `bunsen::ops` is a library of stateless +`bunsen::blocks` is a library of `burn::module::Module` building blocks (stateful layers with trainable parameters), and +`bunsen::ops` is a library of stateless `Tensor` operations. A survey of what's available: ```text @@ -247,8 +240,7 @@ See [`examples/`](examples/) for the full index. At a glance: # Motivation -This library is a synthesis of the utility and extension work that -I've been accumulating in: +This library is a synthesis of the utility and extension work that I've been accumulating in: * * @@ -256,36 +248,29 @@ I've been accumulating in: * * -This library is a work in progress, and I'm working to fold the various -utilities and support code from these projects into a single place; where we -can closely track the burn release cycle, and minimize the dependency-hell -churn problem for writing extensions. +This library is a work in progress, and I'm working to fold the various utilities and support code from these projects +into a single place; where we can closely track the burn release cycle, and minimize the dependency-hell churn problem +for writing extensions. -I plan on continuing to work on this library, and recruit community -involvement for landing and publishing new operators and blocks in a place -we can lock down their testings and documentation. +I plan on continuing to work on this library, and recruit community involvement for landing and publishing new operators +and blocks in a place we can lock down their testings and documentation. ## Future Components -The base libraries have significant features which haven't been polished and stabilized for bunsen -yet. +The base libraries have significant features which haven't been polished and stabilized for bunsen yet. -* weight/data download disk cache - there are several implementations of this in my codebase so far, - the most robust is probably in the `wordchipper` code. -* shard fetching - being able to bind a family of shards to URL template + range pattern; - with information on the target format; and wire that smoothly into the download and cache layer. - this is also currently in some of the LLM/chat codebases. +* weight/data download disk cache - there are several implementations of this in my codebase so far, the most robust is + probably in the `wordchipper` code. +* shard fetching - being able to bind a family of shards to URL template + range pattern; with information on the target + format; and wire that smoothly into the download and cache layer. this is also currently in some of the LLM/chat + codebases. * LLM `DataLoader` - a high-performance burn data loader for LLM models, built on parquet/arrow; and - `wordchipper`. - This is currently in the `zsl-chat` codebase. -* `clap` tooling - I've built a lot of burn-related clap tools, and I'm pretty sure some of the - arguments/setup - machinery - could be shared. + `wordchipper`. This is currently in the `zsl-chat` codebase. +* `clap` tooling - I've built a lot of burn-related clap tools, and I'm pretty sure some of the arguments/setup + machinery could be shared. # License -`bunsen` is distributed under the terms of both the MIT license and the Apache License -(Version 2.0). -See [LICENSE-APACHE](LICENSE-APACHE) and [LICENSE-MIT](LICENSE-MIT) for details. Opening a pull -request is assumed to signal agreement with these licensing terms +`bunsen` is distributed under the terms of both the MIT license and the Apache License (Version 2.0). +See [LICENSE-APACHE](LICENSE-APACHE) and [LICENSE-MIT](LICENSE-MIT) for details. Opening a pull request is assumed to +signal agreement with these licensing terms diff --git a/crates/bunsen/README.md b/crates/bunsen/README.md index 8592d3e..ea609ee 100644 --- a/crates/bunsen/README.md +++ b/crates/bunsen/README.md @@ -8,8 +8,8 @@ [![Discord](https://img.shields.io/discord/1475229838754316502?label=discord)](https://discord.gg/vBgXHWCeah) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zspacelabs/bunsen) -`bunsen` aims to be a "batteries included" complementary -community standard library for extending the [burn](https://burn.dev) tensor library. +`bunsen` aims to be a "batteries included" complementary community standard library for extending +the [burn](https://burn.dev) tensor library. # Book @@ -19,8 +19,8 @@ Read the [bunsen book](https://zspacelabs.ai/bunsen/book) ### Burn Extensions -* `bunsen::burner` - this is a library of `burn::module::Module` lifecycle - components that extend the current functionality of burn. +* `bunsen::burner` - this is a library of `burn::module::Module` lifecycle components that extend the current + functionality of burn. * `bunsen::burner::module::reflection` has powerful tools for dynamic `burn::module::Module` reflection. * `bunsen::burner::optim` has parameter-group optimizer extensions. * `bunsen::burner::tensor` has `Tensor` extension traits (`swap`/`release`, @@ -30,9 +30,8 @@ Read the [bunsen book](https://zspacelabs.ai/bunsen/book) ### Component Libraries -* `bunsen::blocks` - this is a library of `burn::module::Module` components. - This includes simple inner layers, recurrent utility blocks, and entire - model families. +* `bunsen::blocks` - this is a library of `burn::module::Module` components. This includes simple inner layers, + recurrent utility blocks, and entire model families. * `bunsen::ops` - this is a library `burn::tensor::Tensor` operations. * `bunsen::kits` - this is a library of full models and simulation kits. * `bimm` - image models: `resnet`, `swinn` @@ -43,8 +42,8 @@ Read the [bunsen book](https://zspacelabs.ai/bunsen/book) ### App and Testing Support Libs * `bunsen::errors` - this is a library of error types and tooling. -* `bunsen::support` - this is a library of support functions for bunsen, including - testing tooling which may be useful for clients. +* `bunsen::support` - this is a library of support functions for bunsen, including testing tooling which may be useful + for clients. * `bunsen::zspace` - this is a library of z-space / index utilities. # API Examples @@ -55,11 +54,10 @@ A "good parts" survey of some of `bunsen`'s features. See the ## Shape Contracts -`bunsen::contracts` provides allocation-free, always-on runtime tensor-shape -contracts. A contract pairs paper-style shape notation with runtime checking: -it asserts that a tensor's shape matches a declared pattern *and* unpacks named -dimensions for downstream use, catching shape errors at their source. Single -checks run in ~160 ns; the amortized periodic variants average a few ns. +`bunsen::contracts` provides allocation-free, always-on runtime tensor-shape contracts. A contract pairs paper-style +shape notation with runtime checking: +it asserts that a tensor's shape matches a declared pattern *and* unpacks named dimensions for downstream use, catching +shape errors at their source. Single checks run in ~160 ns; the amortized periodic variants average a few ns. ```rust use bunsen::contracts::*; @@ -84,8 +82,8 @@ let [b, h_wins, w_wins, c] = unpack_shape_contract!( ); ``` -In hot loops, use `assert_shape_contract_periodically!` to amortize the check -via exponential backoff while still catching regressions: +In hot loops, use `assert_shape_contract_periodically!` to amortize the check via exponential backoff while still +catching regressions: ```rust use bunsen::contracts::*; @@ -99,15 +97,14 @@ assert_shape_contract_periodically!( ## Tensor Op Extensions -`bunsen::burner::tensor` provides extension traits that add utility methods -directly to `burn::Tensor` — in scope after `use bunsen::burner::tensor::*;`: +`bunsen::burner::tensor` provides extension traits that add utility methods directly to `burn::Tensor` — in scope after +`use bunsen::burner::tensor::*;`: * `TensorOpExt` (all tensor kinds) — `swap` exchanges two tensors in place; `release` moves a tensor out of a field, leaving an empty tensor behind; - `select_dim` selects one index along a dimension and squeezes it, dropping - the rank by one. -* `TensorIntOpExt` (`Int` tensors) — `square`, and `bounded_elem` for - elementwise `[start, end)` range checks producing `Bool` masks. + `select_dim` selects one index along a dimension and squeezes it, dropping the rank by one. +* `TensorIntOpExt` (`Int` tensors) — `square`, and `bounded_elem` for elementwise `[start, end)` range checks producing + `Bool` masks. * `TensorBoolOpExt` (`Bool` tensors) — `count_dim` / `count_dims` count `true` elements along one or more dimensions, with negative indexing. @@ -128,10 +125,9 @@ fn row_counts(grid: Tensor) -> Tensor { ## TensorData Index Views -`TensorDataIndexView` and `TensorDataIndexMutView` wrap burn's low-level -`TensorData` to give ergonomic multi-dimensional element access via bracket -notation — `view[&[i, j]]` — instead of manually flattening indices. The views -deref to the underlying `TensorData`, so `.shape` and friends are right there. +`TensorDataView` and `TensorDataViewMut` wrap burn's low-level +`TensorData` to give ergonomic multi-dimensional element access via bracket notation — `view[&[i, j]]` — instead of +manually flattening indices. The views deref to the underlying `TensorData`, so `.shape` and friends are right there. Handy for inspecting or patching raw tensor data without building full tensors. ```rust @@ -164,10 +160,9 @@ assert_eq!(view[&[0, 0]], 10.0); ## XML Module Reflection `bunsen::burner::module::reflection::XmlModuleTree` turns any burn `Module` -into a queryable XML meta-description of its structure. This enables -type-erased introspection and XPath-style parameter selection — e.g. "every -rank-2 weight under the transformer blocks" — which is exactly what you need to -slice a model into parameter groups for per-group optimizers. +into a queryable XML meta-description of its structure. This enables type-erased introspection and XPath-style parameter +selection — e.g. "every rank-2 weight under the transformer blocks" — which is exactly what you need to slice a model +into parameter groups for per-group optimizers. Take a small container module: @@ -206,8 +201,8 @@ let matrix_params = mtree .to_param_ids() ?; ``` -The dumped structure mirrors the module's fields, with each `@name` taken from -the struct field and a stable `param_id` per tensor: +The dumped structure mirrors the module's fields, with each `@name` taken from the struct field and a stable `param_id` +per tensor: ```xml @@ -233,8 +228,8 @@ the struct field and a stable `param_id` per tensor: ## Blocks & Ops -`bunsen::blocks` is a library of `burn::module::Module` building blocks (stateful -layers with trainable parameters), and `bunsen::ops` is a library of stateless +`bunsen::blocks` is a library of `burn::module::Module` building blocks (stateful layers with trainable parameters), and +`bunsen::ops` is a library of stateless `Tensor` operations. A survey of what's available: ```text @@ -271,5 +266,4 @@ See [`examples/`](https://github.com/zspacelabs/bunsen/tree/main/examples/) for # License -`bunsen` is distributed under the terms of both the MIT license and the Apache License -(Version 2.0). +`bunsen` is distributed under the terms of both the MIT license and the Apache License (Version 2.0). diff --git a/crates/bunsen/src/burner/mod.rs b/crates/bunsen/src/burner/mod.rs index 1bb9a4b..b635fb8 100644 --- a/crates/bunsen/src/burner/mod.rs +++ b/crates/bunsen/src/burner/mod.rs @@ -67,8 +67,8 @@ //! dimensions (negative indexing supported). //! //! [`tensor`] also carries the -//! [`TensorDataIndexView`](tensor::TensorDataIndexView) -//! / [`TensorDataIndexMutView`](tensor::TensorDataIndexMutView) wrappers, +//! [`TensorDataIndexView`](tensor::TensorDataView) +//! / [`TensorDataIndexMutView`](tensor::TensorDataViewMut) wrappers, //! which give `view[&[i, j]]` multi-dimensional element access to a raw //! `TensorData`. diff --git a/crates/bunsen/src/burner/tensor/data_view.rs b/crates/bunsen/src/burner/tensor/data_view.rs index 7681071..1068867 100644 --- a/crates/bunsen/src/burner/tensor/data_view.rs +++ b/crates/bunsen/src/burner/tensor/data_view.rs @@ -11,25 +11,179 @@ use burn::{ prelude::TensorData, tensor::{ AsIndex, + BoolStore, + DType, + DataError, Element, }, }; use crate::zspace::ravel_dims; +/// Extension trait for [`TensorData`] which provides index view builders. +pub trait TensorDataViewExt { + /// Returns an [`Index`] view wrapper of the [`TensorData`]. + /// + /// The view implements [`Deref`]. + /// + /// # Example + /// ```rust,no_run + /// use bunsen::burner::tensor::*; + /// use burn::prelude::*; + /// + /// let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + /// let shape = data.shape.clone(); + /// let view: TensorDataView = data.try_index_view().unwrap(); + /// + /// // Deref + /// assert_eq!(&view.shape, &shape); + /// + /// assert_eq!(view[&[0, 0]], 1.0); + /// assert_eq!(view[&[0, 1]], 2.0); + /// assert_eq!(view[&[1, 0]], 3.0); + /// assert_eq!(view[&[1, 1]], 4.0); + /// ``` + /// + /// # Returns + /// `Ok(view)` on success; `Err(DataError::TypeError)` on view [`DType`] + /// missmatch. + fn try_index_view(&self) -> Result, DataError>; + + /// Returns a [`TensorDataView<'a, E>`] of the [`TensorData`]. + /// + /// The view implements [`Deref`]. + /// + /// # Example + /// ```rust,no_run + /// use bunsen::burner::tensor::*; + /// use burn::prelude::*; + /// + /// let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + /// let shape = data.shape.clone(); + /// let view: TensorDataView = data.expect_index_view(); + /// + /// // Deref + /// assert_eq!(&view.shape, &shape); + /// + /// assert_eq!(view[&[0, 0]], 1.0); + /// assert_eq!(view[&[0, 1]], 2.0); + /// assert_eq!(view[&[1, 0]], 3.0); + /// assert_eq!(view[&[1, 1]], 4.0); + /// ``` + /// + /// # Returns + /// The view. + /// + /// # Panics + /// If the view [`DType`] is not compatible with the [`TensorData`]. + fn expect_index_view(&self) -> TensorDataView<'_, E> { + self.try_index_view().unwrap() + } + + /// Returns a [`TensorDataViewMut<'a, E>`] of the [`TensorData`]. + /// + /// The view implements [`DerefMut`]. + /// + /// # Example + /// ```rust,no_run + /// use bunsen::burner::tensor::*; + /// use burn::prelude::*; + /// + /// let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + /// let shape = data.shape.clone(); + /// let mut view: TensorDataViewMut = data.try_index_mut_view().unwrap(); + /// + /// // Deref + /// assert_eq!(&view.shape, &shape); + /// + /// assert_eq!(view[&[0, 0]], 1.0); + /// assert_eq!(view[&[0, 1]], 2.0); + /// assert_eq!(view[&[1, 0]], 3.0); + /// assert_eq!(view[&[1, 1]], 4.0); + /// + /// view[&[0, 0]] = 10.0; + /// assert_eq!(view[&[0, 0]], 10.0); + /// ``` + /// + /// # Returns + /// `Ok(mut view)` on success; `Err(DataError::TypeError)` on view [`DType`] + /// missmatch. + fn try_index_mut_view(&mut self) -> Result, DataError>; + + /// Returns a [`TensorDataViewMut<'a, E>`] of the [`TensorData`]. + /// + /// The view implements [`DerefMut`]. + /// + /// # Example + /// ```rust,no_run + /// use bunsen::burner::tensor::*; + /// use burn::prelude::*; + /// + /// let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + /// let shape = data.shape.clone(); + /// let mut view: TensorDataViewMut = data.expect_index_mut_view(); + /// + /// // Deref + /// assert_eq!(&view.shape, &shape); + /// + /// assert_eq!(view[&[0, 0]], 1.0); + /// assert_eq!(view[&[0, 1]], 2.0); + /// assert_eq!(view[&[1, 0]], 3.0); + /// assert_eq!(view[&[1, 1]], 4.0); + /// + /// view[&[0, 0]] = 10.0; + /// assert_eq!(view[&[0, 0]], 10.0); + /// ``` + /// + /// # Returns + /// The mut view. + /// + /// # Panics + /// If the view [`DType`] is not compatible with the [`TensorData`]. + fn expect_index_mut_view(&mut self) -> TensorDataViewMut<'_, E> { + self.try_index_mut_view().unwrap() + } +} + +/// This is [`TensorData`]'s method, but should be public. +/// TODO: expose [`TensorData::matches_target_dtype`] upstream. +fn matches_target_dtype(data: &TensorData) -> bool { + let target_dtype = E::dtype(); + match data.dtype { + DType::Bool(BoolStore::U8) => { + matches!(target_dtype, DType::U8 | DType::Bool(BoolStore::U8)) + } + DType::Bool(BoolStore::U32) => { + matches!(target_dtype, DType::U32 | DType::Bool(BoolStore::U32)) + } + dtype => dtype == target_dtype, + } +} + +impl TensorDataViewExt for TensorData { + fn try_index_view(&self) -> Result, DataError> { + TensorDataView::try_view(self) + } + + fn try_index_mut_view(&mut self) -> Result, DataError> { + TensorDataViewMut::try_mut_view(self) + } +} + /// [Index] view wrapper for a [`TensorData`]. /// +/// The view implements [`Deref`]. +/// /// # Example /// ```rust,no_run /// use bunsen::burner::tensor::*; /// use burn::prelude::*; /// /// let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); -/// let shape = data.shape.clone(); -/// let view: TensorDataIndexView = TensorDataIndexView::view(&data); +/// let view: TensorDataView = data.expect_index_view(); /// /// // Deref -/// assert_eq!(&view.shape, &shape); +/// assert_eq!(&view.shape, &data.shape); /// /// assert_eq!(view[&[0, 0]], 1.0); /// assert_eq!(view[&[0, 1]], 2.0); @@ -37,12 +191,12 @@ use crate::zspace::ravel_dims; /// assert_eq!(view[&[1, 1]], 4.0); /// ``` #[derive(Debug)] -pub struct TensorDataIndexView<'a, E: Element> { +pub struct TensorDataView<'a, E: Element> { data: &'a TensorData, _phantom: std::marker::PhantomData<&'a E>, } -impl<'a, E: Element> Deref for TensorDataIndexView<'a, E> { +impl<'a, E: Element> Deref for TensorDataView<'a, E> { type Target = TensorData; fn deref(&self) -> &Self::Target { @@ -50,12 +204,43 @@ impl<'a, E: Element> Deref for TensorDataIndexView<'a, E> { } } -impl<'a, E: Element> TensorDataIndexView<'a, E> { - /// Returns an indexed view of the data. - pub fn view(data: &'a TensorData) -> TensorDataIndexView<'a, E> { - TensorDataIndexView { - data, - _phantom: std::marker::PhantomData, +impl<'a, E: Element> TensorDataView<'a, E> { + /// Returns an [`Index`] view wrapper of the [`TensorData`]. + /// + /// The view implements [`Deref`]. + /// + /// # Example + /// ```rust,no_run + /// use bunsen::burner::tensor::*; + /// use burn::prelude::*; + /// + /// let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + /// let view: TensorDataView = data.try_index_view().unwrap(); + /// + /// // Deref. + /// assert_eq!(&view.shape, &data.shape); + /// + /// assert_eq!(view[&[0, 0]], 1.0); + /// assert_eq!(view[&[0, 1]], 2.0); + /// assert_eq!(view[&[1, 0]], 3.0); + /// assert_eq!(view[&[1, 1]], 4.0); + /// ``` + /// + /// # Returns + /// `Ok(view)` on success; `Err(DataError::TypeError)` on view [`DType`] + /// missmatch. + pub fn try_view(data: &'a TensorData) -> Result, DataError> { + if !matches_target_dtype::(data) { + Err(DataError::TypeMismatch(format!( + "Cannot view TensorData DType \"{:?}\" as \"{:?}\"", + data.dtype, + E::dtype() + ))) + } else { + Ok(TensorDataView { + data, + _phantom: std::marker::PhantomData, + }) } } @@ -68,7 +253,7 @@ impl<'a, E: Element> TensorDataIndexView<'a, E> { } } -impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexView<'a, E> { +impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataView<'a, E> { type Output = E; fn index( @@ -82,6 +267,8 @@ impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexView<'a, E> { /// Mutable [`IndexMut`] view wrapper for a [`TensorData`]. /// +/// The view implements [`DerefMut`]. +/// /// # Example /// ```rust,no_run /// use bunsen::burner::tensor::*; @@ -89,10 +276,10 @@ impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexView<'a, E> { /// /// let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); /// let shape = data.shape.clone(); -/// let mut view: TensorDataIndexMutView = -/// TensorDataIndexMutView::view(&mut data); +/// let mut view: TensorDataViewMut = +/// TensorDataViewMut::try_mut_view(&mut data).unwrap(); /// -/// // Deref +/// // The view implements [`Deref`]. /// assert_eq!(&view.shape, &shape); /// /// assert_eq!(view[&[0, 0]], 1.0); @@ -104,12 +291,12 @@ impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexView<'a, E> { /// assert_eq!(view[&[0, 0]], 10.0); /// ``` #[derive(Debug)] -pub struct TensorDataIndexMutView<'a, E: Element> { +pub struct TensorDataViewMut<'a, E: Element> { data: &'a mut TensorData, _phantom: std::marker::PhantomData<&'a E>, } -impl<'a, E: Element> Deref for TensorDataIndexMutView<'a, E> { +impl<'a, E: Element> Deref for TensorDataViewMut<'a, E> { type Target = TensorData; fn deref(&self) -> &Self::Target { @@ -117,18 +304,54 @@ impl<'a, E: Element> Deref for TensorDataIndexMutView<'a, E> { } } -impl<'a, E: Element> DerefMut for TensorDataIndexMutView<'a, E> { +impl<'a, E: Element> DerefMut for TensorDataViewMut<'a, E> { fn deref_mut(&mut self) -> &mut Self::Target { self.data } } -impl<'a, E: Element> TensorDataIndexMutView<'a, E> { +impl<'a, E: Element> TensorDataViewMut<'a, E> { /// Returns an indexed view of the data. - pub fn view(data: &'a mut TensorData) -> TensorDataIndexMutView<'a, E> { - TensorDataIndexMutView { - data, - _phantom: std::marker::PhantomData, + /// + /// The view implements [`DerefMut`]. + /// + /// # Example + /// ```rust,no_run + /// use bunsen::burner::tensor::*; + /// use burn::prelude::*; + /// + /// let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + /// let shape = data.shape.clone(); + /// let mut view: TensorDataViewMut = + /// TensorDataViewMut::try_mut_view(&mut data).unwrap(); + /// + /// // Deref + /// assert_eq!(&view.shape, &shape); + /// + /// assert_eq!(view[&[0, 0]], 1.0); + /// assert_eq!(view[&[0, 1]], 2.0); + /// assert_eq!(view[&[1, 0]], 3.0); + /// assert_eq!(view[&[1, 1]], 4.0); + /// + /// view[&[0, 0]] = 10.0; + /// assert_eq!(view[&[0, 0]], 10.0); + /// ``` + /// + /// # Returns + /// `Ok(mut view)` on success; `Err(DataError::TypeError)` on view [`DType`] + /// missmatch. + pub fn try_mut_view(data: &'a mut TensorData) -> Result, DataError> { + if !matches_target_dtype::(data) { + Err(DataError::TypeMismatch(format!( + "Cannot view TensorData DType \"{:?}\" as \"{:?}\"", + data.dtype, + E::dtype() + ))) + } else { + Ok(TensorDataViewMut { + data, + _phantom: std::marker::PhantomData, + }) } } @@ -141,7 +364,7 @@ impl<'a, E: Element> TensorDataIndexMutView<'a, E> { } } -impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexMutView<'a, E> { +impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataViewMut<'a, E> { type Output = E; fn index( @@ -153,7 +376,7 @@ impl<'a, I: AsIndex, E: Element> Index<&[I]> for TensorDataIndexMutView<'a, E> { } } -impl<'a, I: AsIndex, E: Element> IndexMut<&[I]> for TensorDataIndexMutView<'a, E> { +impl<'a, I: AsIndex, E: Element> IndexMut<&[I]> for TensorDataViewMut<'a, E> { fn index_mut( &mut self, index: &[I], @@ -167,10 +390,40 @@ impl<'a, I: AsIndex, E: Element> IndexMut<&[I]> for TensorDataIndexMutView<'a, E mod tests { use super::*; + #[test] + fn test_tensor_data_try_view_dtype_mismatch() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + + let result = data.try_index_view::(); + assert!(matches!(result, Err(DataError::TypeMismatch(_)))); + } + + #[test] + #[should_panic(expected = "Cannot view TensorData DType")] + fn test_tensor_data_expect_view_dtype_mismatch() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let _view = data.expect_index_view::(); + } + + #[test] + fn test_tensor_data_try_mut_view_dtype_mismatch() { + let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + + let result = data.try_index_mut_view::(); + assert!(matches!(result, Err(DataError::TypeMismatch(_)))); + } + + #[test] + #[should_panic(expected = "Cannot view TensorData DType")] + fn test_tensor_data_expect_mut_view_dtype_mismatch() { + let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let _view = data.expect_index_mut_view::(); + } + #[test] fn test_tensor_data_index_view() { let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); - let view: TensorDataIndexView = TensorDataIndexView::view(&data); + let view = data.expect_index_view::(); // Deref assert_eq!(&data.shape, &view.shape); @@ -186,7 +439,7 @@ mod tests { let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); let shape = data.shape.clone(); - let mut view: TensorDataIndexMutView = TensorDataIndexMutView::view(&mut data); + let mut view = data.expect_index_mut_view::(); // Deref assert_eq!(&view.shape, &shape); diff --git a/crates/bunsen/src/burner/tensor/mod.rs b/crates/bunsen/src/burner/tensor/mod.rs index 4c8193f..24848ab 100644 --- a/crates/bunsen/src/burner/tensor/mod.rs +++ b/crates/bunsen/src/burner/tensor/mod.rs @@ -45,7 +45,7 @@ //! //! ## `TensorData` Views //! -//! [`TensorDataIndexView`] and [`TensorDataIndexMutView`] wrap a +//! [`TensorDataView`] and [`TensorDataViewMut`] wrap a //! [`burn::tensor::TensorData`] to provide multi-dimensional element access //! via `view[&[i, j]]` indexing. diff --git a/examples/lbm2d_vis/src/main.rs b/examples/lbm2d_vis/src/main.rs index a5e8753..2cac0d9 100644 --- a/examples/lbm2d_vis/src/main.rs +++ b/examples/lbm2d_vis/src/main.rs @@ -14,7 +14,7 @@ use std::{ }; use bunsen::{ - burner::tensor::TensorDataIndexView, + burner::tensor::TensorDataView, kits::sims::lbm::d2q9::{ LBMD2Q9Config, LBMD2Q9State, @@ -25,6 +25,7 @@ use bunsen::{ macroscopic_momentum, }, prelude::{ + TensorDataViewExt, TensorElemOpExt, TensorOpExt, }, @@ -366,10 +367,10 @@ impl FlowVisApp { ) { use graphics::*; - let solid_cells: TensorDataIndexView = TensorDataIndexView::view(&self.solid_mask); + let solid_cells: TensorDataView = self.solid_mask.expect_index_view(); let cell_data = self.get_cell_data(); - let vis_cells: TensorDataIndexView = TensorDataIndexView::view(&cell_data); + let vis_cells: TensorDataView = cell_data.expect_index_view(); let [height, width] = cell_data.shape[0..2].try_into().unwrap(); let [view_width, view_height] = args.viewport().window_size;