torsh_tensor/data_ops/
functions.rs

1//! Auto-generated module
2//!
3//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
4
5use crate::core_ops::Tensor;
6use torsh_core::{
7    device::DeviceType,
8    dtype::TensorElement,
9    error::{Result, TorshError},
10};
11
12impl<T: TensorElement + Copy> Tensor<T> {
13    /// Create tensor from a scalar value repeated to fill the shape
14    pub fn from_scalar(value: T, shape: &[usize], device: DeviceType) -> Result<Self>
15    where
16        T: Copy,
17    {
18        let numel = shape.iter().product::<usize>();
19        let data = vec![value; numel];
20        Self::from_data(data, shape.to_vec(), device)
21    }
22    /// Fill tensor with a single value (in-place)
23    pub fn fill_(&mut self, value: T) -> Result<()>
24    where
25        T: Copy,
26    {
27        for i in 0..self.numel() {
28            self.storage.set(i, value)?;
29        }
30        Ok(())
31    }
32    /// Zero out the tensor (in-place)
33    pub fn zero_(&mut self) -> Result<()>
34    where
35        T: Copy,
36    {
37        self.fill_(T::zero())
38    }
39    /// Fill with ones (in-place)
40    pub fn ones_(&mut self) -> Result<()>
41    where
42        T: Copy,
43    {
44        self.fill_(T::one())
45    }
46    /// Copy data from another tensor (in-place)
47    pub fn copy_(&mut self, other: &Self) -> Result<()>
48    where
49        T: Copy,
50    {
51        if self.shape() != other.shape() {
52            return Err(TorshError::ShapeMismatch {
53                expected: self.shape().dims().to_vec(),
54                got: other.shape().dims().to_vec(),
55            });
56        }
57        let other_data = other.to_vec()?;
58        for (i, &value) in other_data.iter().enumerate() {
59            self.storage.set(i, value)?;
60        }
61        Ok(())
62    }
63    /// Get an element by multi-dimensional index
64    pub fn get_item(&self, indices: &[usize]) -> Result<T>
65    where
66        T: Copy,
67    {
68        if indices.len() != self.ndim() {
69            return Err(TorshError::InvalidArgument(format!(
70                "Expected {} indices, got {}",
71                self.ndim(),
72                indices.len()
73            )));
74        }
75        let binding = self.shape();
76        let shape = binding.dims();
77        for (i, &idx) in indices.iter().enumerate() {
78            if idx >= shape[i] {
79                return Err(TorshError::IndexOutOfBounds {
80                    index: idx,
81                    size: shape[i],
82                });
83            }
84        }
85        let flat_index = self.multi_to_flat_index(indices)?;
86        self.get_item_flat(flat_index)
87    }
88    /// Set an element by multi-dimensional index
89    pub fn set_item(&mut self, indices: &[usize], value: T) -> Result<()>
90    where
91        T: Copy,
92    {
93        if indices.len() != self.ndim() {
94            return Err(TorshError::InvalidArgument(format!(
95                "Expected {} indices, got {}",
96                self.ndim(),
97                indices.len()
98            )));
99        }
100        let binding = self.shape();
101        let shape = binding.dims();
102        for (i, &idx) in indices.iter().enumerate() {
103            if idx >= shape[i] {
104                return Err(TorshError::IndexOutOfBounds {
105                    index: idx,
106                    size: shape[i],
107                });
108            }
109        }
110        let flat_index = self.multi_to_flat_index(indices)?;
111        self.set_item_flat(flat_index, value)
112    }
113    /// Get element by flat index
114    pub fn get_item_flat(&self, index: usize) -> Result<T>
115    where
116        T: Copy,
117    {
118        if index >= self.numel() {
119            return Err(TorshError::IndexOutOfBounds {
120                index,
121                size: self.numel(),
122            });
123        }
124        // Convert flat view-space index → multi-dim → storage index via strides+offset
125        let shape = self.shape();
126        let dims = shape.dims();
127        let ndim = dims.len();
128        let mut multi_dim = vec![0usize; ndim];
129        let mut remaining = index;
130        for i in (0..ndim).rev() {
131            if dims[i] > 0 {
132                multi_dim[i] = remaining % dims[i];
133                remaining /= dims[i];
134            }
135        }
136        let strides = self.strides();
137        let storage_idx = self.storage_offset
138            + multi_dim
139                .iter()
140                .zip(strides.iter())
141                .map(|(&m, &s)| m * s)
142                .sum::<usize>();
143        self.storage.get(storage_idx)
144    }
145    /// Set element by flat index
146    pub fn set_item_flat(&mut self, index: usize, value: T) -> Result<()>
147    where
148        T: Copy,
149    {
150        if index >= self.numel() {
151            return Err(TorshError::IndexOutOfBounds {
152                index,
153                size: self.numel(),
154            });
155        }
156        // Convert flat view-space index → multi-dim → storage index via strides+offset
157        let shape = self.shape();
158        let dims = shape.dims();
159        let ndim = dims.len();
160        let mut multi_dim = vec![0usize; ndim];
161        let mut remaining = index;
162        for i in (0..ndim).rev() {
163            if dims[i] > 0 {
164                multi_dim[i] = remaining % dims[i];
165                remaining /= dims[i];
166            }
167        }
168        let strides = self.strides();
169        let storage_idx = self.storage_offset
170            + multi_dim
171                .iter()
172                .zip(strides.iter())
173                .map(|(&m, &s)| m * s)
174                .sum::<usize>();
175        self.storage.set(storage_idx, value)
176    }
177    /// Convert multi-dimensional indices to flat index
178    pub fn multi_to_flat_index(&self, indices: &[usize]) -> Result<usize> {
179        let binding = self.shape();
180        let shape = binding.dims();
181        if indices.len() != shape.len() {
182            return Err(TorshError::InvalidArgument(format!(
183                "Expected {} indices, got {}",
184                shape.len(),
185                indices.len()
186            )));
187        }
188        let mut flat_index = 0;
189        let mut stride = 1;
190        for i in (0..indices.len()).rev() {
191            flat_index += indices[i] * stride;
192            stride *= shape[i];
193        }
194        Ok(flat_index)
195    }
196    /// Gather values along an axis using indices
197    pub fn gather(&self, dim: usize, indices: &Tensor<i64>) -> Result<Self> {
198        if dim >= self.ndim() {
199            return Err(TorshError::InvalidArgument(format!(
200                "Dimension {} out of range for tensor with {} dimensions",
201                dim,
202                self.ndim()
203            )));
204        }
205        let self_data = self.to_vec()?;
206        let indices_data = indices.to_vec()?;
207        let mut result_data = Vec::new();
208        let result_shape = indices.shape().dims().to_vec();
209        if self.ndim() == 1 {
210            for &index in &indices_data {
211                let idx = if index < 0 {
212                    (self.shape().dims()[0] as i64 + index) as usize
213                } else {
214                    index as usize
215                };
216                if idx >= self.shape().dims()[0] {
217                    return Err(TorshError::InvalidArgument(format!(
218                        "Index {} out of range for tensor with size {}",
219                        index,
220                        self.shape().dims()[0]
221                    )));
222                }
223                result_data.push(self_data[idx]);
224            }
225        } else {
226            let self_shape_ref = self.shape();
227            let self_shape = self_shape_ref.dims();
228            let indices_shape_ref = indices.shape();
229            let indices_shape = indices_shape_ref.dims();
230            let dim_size = self_shape[dim];
231            let mut self_strides = vec![1; self_shape.len()];
232            let mut indices_strides = vec![1; indices_shape.len()];
233            for i in (0..self_shape.len() - 1).rev() {
234                self_strides[i] = self_strides[i + 1] * self_shape[i + 1];
235            }
236            for i in (0..indices_shape.len() - 1).rev() {
237                indices_strides[i] = indices_strides[i + 1] * indices_shape[i + 1];
238            }
239            let total_elements = indices_data.len();
240            for (i, &index_value) in indices_data.iter().enumerate().take(total_elements) {
241                let mut indices_coords = vec![0; indices_shape.len()];
242                let mut temp_i = i;
243                for j in 0..indices_shape.len() {
244                    indices_coords[j] = temp_i / indices_strides[j];
245                    temp_i %= indices_strides[j];
246                }
247                let idx = if index_value < 0 {
248                    (dim_size as i64 + index_value) as usize
249                } else {
250                    index_value as usize
251                };
252                if idx >= dim_size {
253                    return Err(TorshError::InvalidArgument(format!(
254                        "Index {index_value} out of range for dimension {dim} with size {dim_size}"
255                    )));
256                }
257                let mut self_coords = indices_coords.clone();
258                if dim < self_coords.len() {
259                    self_coords[dim] = idx;
260                }
261                let mut flat_idx = 0;
262                for j in 0..self_coords.len() {
263                    flat_idx += self_coords[j] * self_strides[j];
264                }
265                result_data.push(self_data[flat_idx]);
266            }
267        }
268        Self::from_data(result_data, result_shape, self.device)
269    }
270    /// Scatter values along an axis using indices
271    pub fn scatter(&self, dim: usize, indices: &Tensor<i64>, src: &Tensor<T>) -> Result<Self> {
272        if dim >= self.ndim() {
273            return Err(TorshError::InvalidArgument(format!(
274                "Dimension {} out of range for tensor with {} dimensions",
275                dim,
276                self.ndim()
277            )));
278        }
279        let mut result_data = self.to_vec()?;
280        let indices_data = indices.to_vec()?;
281        let src_data = src.to_vec()?;
282        if indices_data.len() != src_data.len() {
283            return Err(TorshError::InvalidArgument(
284                "Indices and source tensor must have the same number of elements".to_string(),
285            ));
286        }
287        if self.ndim() == 1 {
288            for (i, &index) in indices_data.iter().enumerate() {
289                let idx = if index < 0 {
290                    (self.shape().dims()[0] as i64 + index) as usize
291                } else {
292                    index as usize
293                };
294                if idx >= self.shape().dims()[0] {
295                    return Err(TorshError::InvalidArgument(format!(
296                        "Index {} out of range for tensor with size {}",
297                        index,
298                        self.shape().dims()[0]
299                    )));
300                }
301                result_data[idx] = src_data[i];
302            }
303        } else {
304            let self_shape_ref = self.shape();
305            let self_shape = self_shape_ref.dims();
306            let indices_shape_ref = indices.shape();
307            let indices_shape = indices_shape_ref.dims();
308            let dim_size = self_shape[dim];
309            let mut self_strides = vec![1; self_shape.len()];
310            let mut indices_strides = vec![1; indices_shape.len()];
311            for i in (0..self_shape.len() - 1).rev() {
312                self_strides[i] = self_strides[i + 1] * self_shape[i + 1];
313            }
314            for i in (0..indices_shape.len() - 1).rev() {
315                indices_strides[i] = indices_strides[i + 1] * indices_shape[i + 1];
316            }
317            let total_elements = indices_data.len();
318            for (i, &index_value) in indices_data.iter().enumerate().take(total_elements) {
319                let mut indices_coords = vec![0; indices_shape.len()];
320                let mut temp_i = i;
321                for j in 0..indices_shape.len() {
322                    indices_coords[j] = temp_i / indices_strides[j];
323                    temp_i %= indices_strides[j];
324                }
325                let idx = if index_value < 0 {
326                    (dim_size as i64 + index_value) as usize
327                } else {
328                    index_value as usize
329                };
330                if idx >= dim_size {
331                    return Err(TorshError::InvalidArgument(format!(
332                        "Index {index_value} out of range for dimension {dim} with size {dim_size}"
333                    )));
334                }
335                let mut self_coords = indices_coords.clone();
336                if dim < self_coords.len() {
337                    self_coords[dim] = idx;
338                }
339                let mut flat_idx = 0;
340                for j in 0..self_coords.len() {
341                    flat_idx += self_coords[j] * self_strides[j];
342                }
343                result_data[flat_idx] = src_data[i];
344            }
345        }
346        Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
347    }
348    /// Scatter values along an axis using indices and add to existing values
349    ///
350    /// # PyTorch Compatibility
351    /// Equivalent to `torch.scatter_add(tensor, dim, index, src)`
352    ///
353    /// # Arguments
354    /// * `dim` - Dimension along which to index
355    /// * `indices` - Index tensor (same shape as src)
356    /// * `src` - Source tensor containing values to add
357    ///
358    /// # Examples
359    /// ```ignore
360    /// let tensor = Tensor::zeros(&[5], DeviceType::Cpu)?;
361    /// let indices = Tensor::from_data(vec![0i64, 1, 2, 0, 1], vec![5], DeviceType::Cpu)?;
362    /// let src = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0, 5.0], vec![5], DeviceType::Cpu)?;
363    /// let result = tensor.scatter_add(0, &indices, &src)?;
364    /// // result[0] += 1.0 + 4.0 = 5.0
365    /// // result[1] += 2.0 + 5.0 = 7.0
366    /// // result[2] += 3.0 = 3.0
367    /// ```
368    pub fn scatter_add(&self, dim: usize, indices: &Tensor<i64>, src: &Tensor<T>) -> Result<Self>
369    where
370        T: std::ops::Add<Output = T>,
371    {
372        if dim >= self.ndim() {
373            return Err(TorshError::InvalidArgument(format!(
374                "Dimension {} out of range for tensor with {} dimensions",
375                dim,
376                self.ndim()
377            )));
378        }
379        let mut result_data = self.to_vec()?;
380        let indices_data = indices.to_vec()?;
381        let src_data = src.to_vec()?;
382        if indices_data.len() != src_data.len() {
383            return Err(TorshError::InvalidArgument(
384                "Indices and source tensor must have the same number of elements".to_string(),
385            ));
386        }
387        if self.ndim() == 1 {
388            for (i, &index) in indices_data.iter().enumerate() {
389                let idx = if index < 0 {
390                    (self.shape().dims()[0] as i64 + index) as usize
391                } else {
392                    index as usize
393                };
394                if idx >= self.shape().dims()[0] {
395                    return Err(TorshError::InvalidArgument(format!(
396                        "Index {} out of range for tensor with size {}",
397                        index,
398                        self.shape().dims()[0]
399                    )));
400                }
401                result_data[idx] = result_data[idx] + src_data[i];
402            }
403        } else {
404            let self_shape_ref = self.shape();
405            let self_shape = self_shape_ref.dims();
406            let indices_shape_ref = indices.shape();
407            let indices_shape = indices_shape_ref.dims();
408            let dim_size = self_shape[dim];
409            let mut self_strides = vec![1; self_shape.len()];
410            let mut indices_strides = vec![1; indices_shape.len()];
411            for i in (0..self_shape.len() - 1).rev() {
412                self_strides[i] = self_strides[i + 1] * self_shape[i + 1];
413            }
414            for i in (0..indices_shape.len() - 1).rev() {
415                indices_strides[i] = indices_strides[i + 1] * indices_shape[i + 1];
416            }
417            let total_elements = indices_data.len();
418            for (i, &index_value) in indices_data.iter().enumerate().take(total_elements) {
419                let mut indices_coords = vec![0; indices_shape.len()];
420                let mut temp_i = i;
421                for j in 0..indices_shape.len() {
422                    indices_coords[j] = temp_i / indices_strides[j];
423                    temp_i %= indices_strides[j];
424                }
425                let idx = if index_value < 0 {
426                    (dim_size as i64 + index_value) as usize
427                } else {
428                    index_value as usize
429                };
430                if idx >= dim_size {
431                    return Err(TorshError::InvalidArgument(format!(
432                        "Index {index_value} out of range for dimension {dim} with size {dim_size}"
433                    )));
434                }
435                let mut self_coords = indices_coords.clone();
436                if dim < self_coords.len() {
437                    self_coords[dim] = idx;
438                }
439                let mut flat_idx = 0;
440                for j in 0..self_coords.len() {
441                    flat_idx += self_coords[j] * self_strides[j];
442                }
443                result_data[flat_idx] = result_data[flat_idx] + src_data[i];
444            }
445        }
446        Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
447    }
448    /// Repeat tensor along specified dimensions
449    pub fn repeat(&self, repeats: &[usize]) -> Result<Self> {
450        if repeats.len() != self.ndim() {
451            return Err(TorshError::InvalidArgument(format!(
452                "Number of repeats {} must match tensor dimensions {}",
453                repeats.len(),
454                self.ndim()
455            )));
456        }
457        let self_data = self.to_vec()?;
458        let shape_binding = self.shape();
459        let self_shape = shape_binding.dims();
460        let new_shape: Vec<usize> = self_shape
461            .iter()
462            .zip(repeats.iter())
463            .map(|(&dim, &repeat)| dim * repeat)
464            .collect();
465        let new_numel = new_shape.iter().product();
466        let mut result_data = Vec::with_capacity(new_numel);
467        for result_idx in 0..new_numel {
468            let mut result_coords = vec![0; new_shape.len()];
469            let mut temp_idx = result_idx;
470            for i in (0..new_shape.len()).rev() {
471                result_coords[i] = temp_idx % new_shape[i];
472                temp_idx /= new_shape[i];
473            }
474            let source_coords: Vec<usize> = result_coords
475                .iter()
476                .zip(self_shape.iter())
477                .map(|(&result_coord, &dim_size)| result_coord % dim_size)
478                .collect();
479            let mut source_idx = 0;
480            let mut stride = 1;
481            for i in (0..self_shape.len()).rev() {
482                source_idx += source_coords[i] * stride;
483                stride *= self_shape[i];
484            }
485            result_data.push(self_data[source_idx]);
486        }
487        Self::from_data(result_data, new_shape, self.device)
488    }
489    /// Add values to tensor at specified indices along a dimension
490    ///
491    /// # PyTorch Compatibility
492    /// Equivalent to `torch.index_add(tensor, dim, index, source, alpha=1.0)`
493    ///
494    /// # Arguments
495    /// * `dim` - Dimension along which to index
496    /// * `index` - 1D tensor containing indices
497    /// * `source` - Source tensor to add
498    ///
499    /// # Examples
500    /// ```ignore
501    /// let tensor = Tensor::zeros(&[3, 5], DeviceType::Cpu)?;
502    /// let index = Tensor::from_data(vec![0i64, 2], vec![2], DeviceType::Cpu)?;
503    /// let source = Tensor::ones(&[2, 5], DeviceType::Cpu)?;
504    /// let result = tensor.index_add(0, &index, &source)?;
505    /// ```
506    pub fn index_add(&self, dim: isize, index: &Tensor<i64>, source: &Self) -> Result<Self>
507    where
508        T: std::ops::Add<Output = T>,
509    {
510        let ndim = self.ndim();
511        let dim = if dim < 0 {
512            (ndim as isize + dim) as usize
513        } else {
514            dim as usize
515        };
516        if dim >= ndim {
517            return Err(TorshError::InvalidArgument(format!(
518                "Dimension {} out of range for {}-D tensor",
519                dim, ndim
520            )));
521        }
522        if index.ndim() != 1 {
523            return Err(TorshError::InvalidArgument(
524                "index must be 1D tensor".to_string(),
525            ));
526        }
527        let index_size = index.shape().dims()[0];
528        let self_shape = self.shape().to_vec();
529        let source_shape = source.shape().to_vec();
530        if source_shape.len() != self_shape.len() {
531            return Err(TorshError::ShapeMismatch {
532                expected: self_shape.clone(),
533                got: source_shape.clone(),
534            });
535        }
536        for (i, (&s, &src_s)) in self_shape.iter().zip(source_shape.iter()).enumerate() {
537            if i == dim {
538                if src_s != index_size {
539                    return Err(TorshError::InvalidArgument(format!(
540                        "source dimension {} size {} must match index size {}",
541                        i, src_s, index_size
542                    )));
543                }
544            } else if s != src_s {
545                return Err(TorshError::ShapeMismatch {
546                    expected: self_shape.clone(),
547                    got: source_shape.clone(),
548                });
549            }
550        }
551        let mut result_data = self.to_vec()?;
552        let source_data = source.to_vec()?;
553        let index_data = index.to_vec()?;
554        let dim_size = self_shape[dim];
555        let outer_size: usize = self_shape[..dim].iter().product();
556        let inner_size: usize = self_shape[dim + 1..].iter().product();
557        for (src_idx_in_dim, &target_idx) in index_data.iter().enumerate() {
558            if target_idx < 0 || target_idx as usize >= dim_size {
559                return Err(TorshError::InvalidArgument(format!(
560                    "Index {} out of range for dimension size {}",
561                    target_idx, dim_size
562                )));
563            }
564            let target_idx = target_idx as usize;
565            for outer in 0..outer_size {
566                for inner in 0..inner_size {
567                    let result_idx =
568                        outer * dim_size * inner_size + target_idx * inner_size + inner;
569                    let source_idx =
570                        outer * index_size * inner_size + src_idx_in_dim * inner_size + inner;
571                    result_data[result_idx] = result_data[result_idx] + source_data[source_idx];
572                }
573            }
574        }
575        Self::from_data(result_data, self_shape, self.device)
576    }
577    /// Copy values from source to tensor at specified indices along a dimension
578    ///
579    /// # PyTorch Compatibility
580    /// Equivalent to `torch.index_copy(tensor, dim, index, source)`
581    ///
582    /// # Arguments
583    /// * `dim` - Dimension along which to index
584    /// * `index` - 1D tensor containing indices
585    /// * `source` - Source tensor to copy from
586    ///
587    /// # Examples
588    /// ```ignore
589    /// let tensor = Tensor::zeros(&[3, 5], DeviceType::Cpu)?;
590    /// let index = Tensor::from_data(vec![0i64, 2], vec![2], DeviceType::Cpu)?;
591    /// let source = Tensor::ones(&[2, 5], DeviceType::Cpu)?;
592    /// let result = tensor.index_copy(0, &index, &source)?;
593    /// ```
594    pub fn index_copy(&self, dim: isize, index: &Tensor<i64>, source: &Self) -> Result<Self> {
595        let ndim = self.ndim();
596        let dim = if dim < 0 {
597            (ndim as isize + dim) as usize
598        } else {
599            dim as usize
600        };
601        if dim >= ndim {
602            return Err(TorshError::InvalidArgument(format!(
603                "Dimension {} out of range for {}-D tensor",
604                dim, ndim
605            )));
606        }
607        if index.ndim() != 1 {
608            return Err(TorshError::InvalidArgument(
609                "index must be 1D tensor".to_string(),
610            ));
611        }
612        let index_size = index.shape().dims()[0];
613        let self_shape = self.shape().to_vec();
614        let source_shape = source.shape().to_vec();
615        if source_shape.len() != self_shape.len() {
616            return Err(TorshError::ShapeMismatch {
617                expected: self_shape.clone(),
618                got: source_shape.clone(),
619            });
620        }
621        for (i, (&s, &src_s)) in self_shape.iter().zip(source_shape.iter()).enumerate() {
622            if i == dim {
623                if src_s != index_size {
624                    return Err(TorshError::InvalidArgument(format!(
625                        "source dimension {} size {} must match index size {}",
626                        i, src_s, index_size
627                    )));
628                }
629            } else if s != src_s {
630                return Err(TorshError::ShapeMismatch {
631                    expected: self_shape.clone(),
632                    got: source_shape.clone(),
633                });
634            }
635        }
636        let mut result_data = self.to_vec()?;
637        let source_data = source.to_vec()?;
638        let index_data = index.to_vec()?;
639        let dim_size = self_shape[dim];
640        let outer_size: usize = self_shape[..dim].iter().product();
641        let inner_size: usize = self_shape[dim + 1..].iter().product();
642        for (src_idx_in_dim, &target_idx) in index_data.iter().enumerate() {
643            if target_idx < 0 || target_idx as usize >= dim_size {
644                return Err(TorshError::InvalidArgument(format!(
645                    "Index {} out of range for dimension size {}",
646                    target_idx, dim_size
647                )));
648            }
649            let target_idx = target_idx as usize;
650            for outer in 0..outer_size {
651                for inner in 0..inner_size {
652                    let result_idx =
653                        outer * dim_size * inner_size + target_idx * inner_size + inner;
654                    let source_idx =
655                        outer * index_size * inner_size + src_idx_in_dim * inner_size + inner;
656                    result_data[result_idx] = source_data[source_idx];
657                }
658            }
659        }
660        Self::from_data(result_data, self_shape, self.device)
661    }
662    /// Fill values in tensor at specified indices along a dimension
663    ///
664    /// # PyTorch Compatibility
665    /// Equivalent to `torch.index_fill(tensor, dim, index, value)`
666    ///
667    /// # Arguments
668    /// * `dim` - Dimension along which to index
669    /// * `index` - 1D tensor containing indices
670    /// * `value` - Scalar value to fill
671    ///
672    /// # Examples
673    /// ```ignore
674    /// let tensor = Tensor::zeros(&[3, 5], DeviceType::Cpu)?;
675    /// let index = Tensor::from_data(vec![0i64, 2], vec![2], DeviceType::Cpu)?;
676    /// let result = tensor.index_fill(0, &index, 3.14)?;
677    /// ```
678    pub fn index_fill(&self, dim: isize, index: &Tensor<i64>, value: T) -> Result<Self> {
679        let ndim = self.ndim();
680        let dim = if dim < 0 {
681            (ndim as isize + dim) as usize
682        } else {
683            dim as usize
684        };
685        if dim >= ndim {
686            return Err(TorshError::InvalidArgument(format!(
687                "Dimension {} out of range for {}-D tensor",
688                dim, ndim
689            )));
690        }
691        if index.ndim() != 1 {
692            return Err(TorshError::InvalidArgument(
693                "index must be 1D tensor".to_string(),
694            ));
695        }
696        let mut result_data = self.to_vec()?;
697        let index_data = index.to_vec()?;
698        let self_shape = self.shape().to_vec();
699        let dim_size = self_shape[dim];
700        let outer_size: usize = self_shape[..dim].iter().product();
701        let inner_size: usize = self_shape[dim + 1..].iter().product();
702        for &target_idx in index_data.iter() {
703            if target_idx < 0 || target_idx as usize >= dim_size {
704                return Err(TorshError::InvalidArgument(format!(
705                    "Index {} out of range for dimension size {}",
706                    target_idx, dim_size
707                )));
708            }
709            let target_idx = target_idx as usize;
710            for outer in 0..outer_size {
711                for inner in 0..inner_size {
712                    let result_idx =
713                        outer * dim_size * inner_size + target_idx * inner_size + inner;
714                    result_data[result_idx] = value;
715                }
716            }
717        }
718        Self::from_data(result_data, self_shape, self.device)
719    }
720    /// Place values at specified flat indices (in-place-like operation, returns new tensor)
721    ///
722    /// # PyTorch Compatibility
723    /// Equivalent to `torch.put_(tensor, indices, values)` but returns new tensor
724    ///
725    /// # Arguments
726    /// * `indices` - 1D tensor of flat indices
727    /// * `values` - 1D tensor of values (must match indices length or be broadcastable)
728    ///
729    /// # Examples
730    /// ```ignore
731    /// let tensor = Tensor::zeros(&[3, 3], DeviceType::Cpu)?;  // [[0,0,0],[0,0,0],[0,0,0]]
732    /// let indices = Tensor::from_data(vec![0i64, 4, 8], vec![3], DeviceType::Cpu)?;
733    /// let values = Tensor::from_data(vec![1.0f32, 2.0, 3.0], vec![3], DeviceType::Cpu)?;
734    /// let result = tensor.put_(&indices, &values)?;  // [[1,0,0],[0,2,0],[0,0,3]]
735    /// ```
736    pub fn put_(&self, indices: &Tensor<i64>, values: &Tensor<T>) -> Result<Self> {
737        if indices.ndim() != 1 {
738            return Err(TorshError::InvalidArgument(
739                "indices must be 1D tensor".to_string(),
740            ));
741        }
742        if values.ndim() != 1 {
743            return Err(TorshError::InvalidArgument(
744                "values must be 1D tensor".to_string(),
745            ));
746        }
747        let indices_data = indices.to_vec()?;
748        let values_data = values.to_vec()?;
749        if indices_data.len() != values_data.len() {
750            return Err(TorshError::InvalidArgument(format!(
751                "Number of values {} must match number of indices {}",
752                values_data.len(),
753                indices_data.len()
754            )));
755        }
756        let mut result_data = self.to_vec()?;
757        let numel = self.numel();
758        for (i, &index) in indices_data.iter().enumerate() {
759            let idx = if index < 0 {
760                ((numel as i64) + index) as usize
761            } else {
762                index as usize
763            };
764            if idx >= numel {
765                return Err(TorshError::InvalidArgument(format!(
766                    "Index {} out of range for tensor with {} elements",
767                    index, numel
768                )));
769            }
770            result_data[idx] = values_data[i];
771        }
772        Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
773    }
774    /// Scatter values from source tensor where mask is true (PyTorch-compatible)
775    ///
776    /// Copies values from the source tensor to positions where the mask is true.
777    /// The mask must have the same shape as self. Source values are taken sequentially
778    /// and placed at positions where mask is true.
779    ///
780    /// # PyTorch Compatibility
781    /// Equivalent to `torch.masked_scatter(tensor, mask, source)`
782    ///
783    /// # Arguments
784    /// * `mask` - Boolean tensor with same shape as self
785    /// * `source` - Tensor containing values to scatter (must have at least as many elements as true values in mask)
786    ///
787    /// # Examples
788    /// ```ignore
789    /// let tensor = Tensor::zeros(&[3, 3], DeviceType::Cpu)?;
790    /// let mask = Tensor::from_data(
791    ///     vec![true, false, false, false, true, false, false, false, true],
792    ///     vec![3, 3],
793    ///     DeviceType::Cpu
794    /// )?;
795    /// let source = Tensor::from_data(vec![1.0f32, 2.0, 3.0], vec![3], DeviceType::Cpu)?;
796    /// let result = tensor.masked_scatter(&mask, &source)?;  // [[1,0,0],[0,2,0],[0,0,3]]
797    /// ```
798    pub fn masked_scatter(&self, mask: &Tensor<bool>, source: &Tensor<T>) -> Result<Self> {
799        if self.shape() != mask.shape() {
800            return Err(TorshError::ShapeMismatch {
801                expected: self.shape().dims().to_vec(),
802                got: mask.shape().dims().to_vec(),
803            });
804        }
805        let mask_data = mask.to_vec()?;
806        let true_count = mask_data.iter().filter(|&&x| x).count();
807        if source.numel() < true_count {
808            return Err(TorshError::InvalidArgument(format!(
809                "Source tensor has {} elements but need {} for scatter (mask has {} true values)",
810                source.numel(),
811                true_count,
812                true_count
813            )));
814        }
815        let self_data = self.to_vec()?;
816        let source_data = source.to_vec()?;
817        let mut result_data = Vec::with_capacity(self_data.len());
818        let mut source_idx = 0;
819        for (i, &self_val) in self_data.iter().enumerate() {
820            if i < mask_data.len() && mask_data[i] {
821                result_data.push(source_data[source_idx]);
822                source_idx += 1;
823            } else {
824                result_data.push(self_val);
825            }
826        }
827        Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
828    }
829    /// Multi-dimensional indexed put operation (PyTorch-compatible)
830    ///
831    /// Places values from source tensor at positions specified by index tensors.
832    /// Each index tensor specifies indices along one dimension. Index tensors must
833    /// be broadcastable to the same shape.
834    ///
835    /// # PyTorch Compatibility
836    /// Equivalent to `torch.index_put(tensor, indices, values)` where indices is a tuple of index tensors
837    ///
838    /// # Arguments
839    /// * `indices` - Slice of index tensors, one per dimension to index
840    /// * `values` - Tensor of values to place (must broadcast to indexed positions)
841    ///
842    /// # Examples
843    /// ```ignore
844    /// // 2D example: index_put a 3x3 matrix with row=[0,1] col=[1,2]
845    /// let tensor = Tensor::zeros(&[3, 3], DeviceType::Cpu)?;
846    /// let row_idx = Tensor::from_data(vec![0i64, 1], vec![2], DeviceType::Cpu)?;
847    /// let col_idx = Tensor::from_data(vec![1i64, 2], vec![2], DeviceType::Cpu)?;
848    /// let values = Tensor::from_data(vec![10.0f32, 20.0], vec![2], DeviceType::Cpu)?;
849    /// let result = tensor.index_put(&[row_idx, col_idx], &values)?;
850    /// // result[0,1] = 10.0, result[1,2] = 20.0
851    /// ```
852    pub fn index_put(&self, indices: &[Tensor<i64>], values: &Tensor<T>) -> Result<Self> {
853        if indices.is_empty() {
854            return Err(TorshError::InvalidArgument(
855                "indices cannot be empty".to_string(),
856            ));
857        }
858        if indices.len() > self.ndim() {
859            return Err(TorshError::InvalidArgument(format!(
860                "Too many indices ({}) for tensor with {} dimensions",
861                indices.len(),
862                self.ndim()
863            )));
864        }
865        let index_shape_ref = indices[0].shape();
866        let index_shape = index_shape_ref.dims();
867        let num_indices = indices[0].numel();
868        for idx_tensor in indices.iter() {
869            if idx_tensor.shape().dims() != index_shape {
870                return Err(TorshError::ShapeMismatch {
871                    expected: index_shape.to_vec(),
872                    got: idx_tensor.shape().dims().to_vec(),
873                });
874            }
875        }
876        if values.numel() != num_indices && values.numel() != 1 {
877            return Err(TorshError::InvalidArgument(format!(
878                "Values tensor has {} elements but need {} (or 1 for broadcasting)",
879                values.numel(),
880                num_indices
881            )));
882        }
883        let mut result_data = self.to_vec()?;
884        let self_shape_ref = self.shape();
885        let self_shape = self_shape_ref.dims();
886        let values_data = values.to_vec()?;
887        let index_data: Result<Vec<Vec<i64>>> = indices.iter().map(|idx| idx.to_vec()).collect();
888        let index_data = index_data?;
889        let mut strides = vec![1; self_shape.len()];
890        for i in (0..self_shape.len() - 1).rev() {
891            strides[i] = strides[i + 1] * self_shape[i + 1];
892        }
893        for i in 0..num_indices {
894            let value = if values_data.len() == 1 {
895                values_data[0]
896            } else {
897                values_data[i]
898            };
899            let mut flat_idx = 0;
900            for (dim, idx_vec) in index_data.iter().enumerate() {
901                let mut idx = idx_vec[i];
902                if idx < 0 {
903                    idx += self_shape[dim] as i64;
904                }
905                if idx < 0 || idx >= self_shape[dim] as i64 {
906                    return Err(TorshError::InvalidArgument(format!(
907                        "Index {} out of bounds for dimension {} with size {}",
908                        idx_vec[i], dim, self_shape[dim]
909                    )));
910                }
911                flat_idx += (idx as usize) * strides[dim];
912            }
913            result_data[flat_idx] = value;
914        }
915        Self::from_data(result_data, self_shape.to_vec(), self.device)
916    }
917    /// Scatter with reduction operation (PyTorch-compatible)
918    ///
919    /// Generalized scatter operation that applies a reduction operation (sum, prod, mean, etc.)
920    /// when scattering values to the same index position.
921    ///
922    /// # PyTorch Compatibility
923    /// Equivalent to `torch.scatter_reduce(tensor, dim, index, src, reduce)`
924    ///
925    /// # Arguments
926    /// * `dim` - Dimension along which to scatter
927    /// * `indices` - Index tensor specifying where to scatter values
928    /// * `src` - Source tensor containing values to scatter
929    /// * `reduce` - Reduction operation ("sum", "prod", "mean", "amax", "amin")
930    ///
931    /// # Examples
932    /// ```ignore
933    /// let tensor = Tensor::zeros(&[5], DeviceType::Cpu)?;
934    /// let indices = Tensor::from_data(vec![0i64, 1, 2, 0, 1], vec![5], DeviceType::Cpu)?;
935    /// let src = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0, 5.0], vec![5], DeviceType::Cpu)?;
936    /// let result = tensor.scatter_reduce(0, &indices, &src, "sum")?;
937    /// // result[0] = 1.0 + 4.0 = 5.0 (sum reduction)
938    /// // result[1] = 2.0 + 5.0 = 7.0
939    /// ```
940    pub fn scatter_reduce(
941        &self,
942        dim: usize,
943        indices: &Tensor<i64>,
944        src: &Tensor<T>,
945        reduce: &str,
946    ) -> Result<Self>
947    where
948        T: std::ops::Add<Output = T>
949            + std::ops::Mul<Output = T>
950            + std::ops::Div<Output = T>
951            + PartialOrd
952            + num_traits::FromPrimitive,
953    {
954        if dim >= self.ndim() {
955            return Err(TorshError::InvalidArgument(format!(
956                "Dimension {} out of range for {}-dimensional tensor",
957                dim,
958                self.ndim()
959            )));
960        }
961        if indices.shape() != src.shape() {
962            return Err(TorshError::ShapeMismatch {
963                expected: indices.shape().dims().to_vec(),
964                got: src.shape().dims().to_vec(),
965            });
966        }
967        let indices_data = indices.to_vec()?;
968        let src_data = src.to_vec()?;
969        let mut result_data = self.to_vec()?;
970        let self_shape_ref = self.shape();
971        let self_shape = self_shape_ref.dims();
972        let mut counts = if reduce == "mean" {
973            vec![0usize; result_data.len()]
974        } else {
975            vec![]
976        };
977        if self.ndim() == 1 {
978            for (i, &index) in indices_data.iter().enumerate() {
979                let idx = if index < 0 {
980                    (self_shape[0] as i64 + index) as usize
981                } else {
982                    index as usize
983                };
984                if idx >= self_shape[0] {
985                    return Err(TorshError::InvalidArgument(format!(
986                        "Index {} out of bounds for dimension size {}",
987                        index, self_shape[0]
988                    )));
989                }
990                result_data[idx] = match reduce {
991                    "sum" => result_data[idx] + src_data[i],
992                    "prod" => result_data[idx] * src_data[i],
993                    "mean" => {
994                        counts[idx] += 1;
995                        result_data[idx] + src_data[i]
996                    }
997                    "amax" => {
998                        if src_data[i] > result_data[idx] {
999                            src_data[i]
1000                        } else {
1001                            result_data[idx]
1002                        }
1003                    }
1004                    "amin" => {
1005                        if src_data[i] < result_data[idx] {
1006                            src_data[i]
1007                        } else {
1008                            result_data[idx]
1009                        }
1010                    }
1011                    _ => {
1012                        return Err(TorshError::InvalidArgument(format!(
1013                            "Unknown reduce operation: {}. Supported: sum, prod, mean, amax, amin",
1014                            reduce
1015                        )));
1016                    }
1017                };
1018            }
1019            if reduce == "mean" {
1020                for (i, count) in counts.iter().enumerate() {
1021                    if *count > 0 {
1022                        result_data[i] = T::from_usize(*count)
1023                            .and_then(|c| Some(result_data[i] / c))
1024                            .unwrap_or(result_data[i]);
1025                    }
1026                }
1027            }
1028        } else {
1029            let dim_size = self_shape[dim];
1030            let _outer_size: usize = self_shape[..dim].iter().product();
1031            let _inner_size: usize = self_shape[dim + 1..].iter().product();
1032            let mut self_strides = vec![1; self_shape.len()];
1033            for i in (0..self_shape.len() - 1).rev() {
1034                self_strides[i] = self_strides[i + 1] * self_shape[i + 1];
1035            }
1036            let src_shape_ref = src.shape();
1037            let src_shape = src_shape_ref.dims();
1038            let mut src_strides = vec![1; src_shape.len()];
1039            for i in (0..src_shape.len() - 1).rev() {
1040                src_strides[i] = src_strides[i + 1] * src_shape[i + 1];
1041            }
1042            for i in 0..indices_data.len() {
1043                let index = indices_data[i];
1044                let idx = if index < 0 {
1045                    (dim_size as i64 + index) as usize
1046                } else {
1047                    index as usize
1048                };
1049                if idx >= dim_size {
1050                    return Err(TorshError::InvalidArgument(format!(
1051                        "Index {} out of bounds for dimension {} size {}",
1052                        index, dim, dim_size
1053                    )));
1054                }
1055                let mut coords = vec![0; self_shape.len()];
1056                let mut remainder = i;
1057                for (d, &stride) in src_strides.iter().enumerate() {
1058                    coords[d] = remainder / stride;
1059                    remainder %= stride;
1060                }
1061                coords[dim] = idx;
1062                let flat_idx = coords
1063                    .iter()
1064                    .zip(self_strides.iter())
1065                    .map(|(c, s)| c * s)
1066                    .sum::<usize>();
1067                result_data[flat_idx] = match reduce {
1068                    "sum" => result_data[flat_idx] + src_data[i],
1069                    "prod" => result_data[flat_idx] * src_data[i],
1070                    "mean" => {
1071                        counts[flat_idx] += 1;
1072                        result_data[flat_idx] + src_data[i]
1073                    }
1074                    "amax" => {
1075                        if src_data[i] > result_data[flat_idx] {
1076                            src_data[i]
1077                        } else {
1078                            result_data[flat_idx]
1079                        }
1080                    }
1081                    "amin" => {
1082                        if src_data[i] < result_data[flat_idx] {
1083                            src_data[i]
1084                        } else {
1085                            result_data[flat_idx]
1086                        }
1087                    }
1088                    _ => {
1089                        return Err(TorshError::InvalidArgument(format!(
1090                            "Unknown reduce operation: {}",
1091                            reduce
1092                        )));
1093                    }
1094                };
1095            }
1096            if reduce == "mean" {
1097                for (i, count) in counts.iter().enumerate() {
1098                    if *count > 0 {
1099                        result_data[i] = T::from_usize(*count)
1100                            .and_then(|c| Some(result_data[i] / c))
1101                            .unwrap_or(result_data[i]);
1102                    }
1103                }
1104            }
1105        }
1106        Self::from_data(result_data, self_shape.to_vec(), self.device)
1107    }
1108    /// Scatter values to the diagonal (PyTorch-compatible)
1109    ///
1110    /// Embeds the values of src tensor into self along the diagonal elements,
1111    /// with respect to dim1 and dim2. The offset determines which diagonal to use.
1112    ///
1113    /// # PyTorch Compatibility
1114    /// Equivalent to `torch.diagonal_scatter(tensor, src, offset, dim1, dim2)`
1115    ///
1116    /// # Arguments
1117    /// * `src` - Source tensor containing values for the diagonal
1118    /// * `offset` - Diagonal offset (0=main diagonal, >0=above, <0=below)
1119    /// * `dim1` - First dimension (default: 0)
1120    /// * `dim2` - Second dimension (default: 1)
1121    ///
1122    /// # Examples
1123    /// ```ignore
1124    /// let tensor = Tensor::zeros(&[3, 3], DeviceType::Cpu)?;
1125    /// let src = Tensor::from_data(vec![1.0f32, 2.0, 3.0], vec![3], DeviceType::Cpu)?;
1126    /// let result = tensor.diagonal_scatter(&src, 0, 0, 1)?;
1127    /// // result = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
1128    /// ```
1129    pub fn diagonal_scatter(
1130        &self,
1131        src: &Tensor<T>,
1132        offset: isize,
1133        dim1: usize,
1134        dim2: usize,
1135    ) -> Result<Self> {
1136        if dim1 >= self.ndim() || dim2 >= self.ndim() {
1137            return Err(TorshError::InvalidArgument(format!(
1138                "Dimensions ({}, {}) out of range for {}-dimensional tensor",
1139                dim1,
1140                dim2,
1141                self.ndim()
1142            )));
1143        }
1144        if dim1 == dim2 {
1145            return Err(TorshError::InvalidArgument(
1146                "dim1 and dim2 must be different".to_string(),
1147            ));
1148        }
1149        let self_shape_ref = self.shape();
1150        let self_shape = self_shape_ref.dims();
1151        let dim1_size = self_shape[dim1];
1152        let dim2_size = self_shape[dim2];
1153        let diag_len = if offset >= 0 {
1154            let offset_u = offset as usize;
1155            if offset_u >= dim2_size {
1156                0
1157            } else {
1158                std::cmp::min(dim1_size, dim2_size - offset_u)
1159            }
1160        } else {
1161            let offset_u = (-offset) as usize;
1162            if offset_u >= dim1_size {
1163                0
1164            } else {
1165                std::cmp::min(dim1_size - offset_u, dim2_size)
1166            }
1167        };
1168        if src.numel() != diag_len {
1169            return Err(TorshError::ShapeMismatch {
1170                expected: vec![diag_len],
1171                got: vec![src.numel()],
1172            });
1173        }
1174        let mut result_data = self.to_vec()?;
1175        let src_data = src.to_vec()?;
1176        let mut strides = vec![1; self_shape.len()];
1177        for i in (0..self_shape.len() - 1).rev() {
1178            strides[i] = strides[i + 1] * self_shape[i + 1];
1179        }
1180        for i in 0..diag_len {
1181            let mut indices = vec![0; self_shape.len()];
1182            if offset >= 0 {
1183                indices[dim1] = i;
1184                indices[dim2] = i + offset as usize;
1185            } else {
1186                indices[dim1] = i + (-offset) as usize;
1187                indices[dim2] = i;
1188            }
1189            let mut flat_idx = 0;
1190            for (d, &idx) in indices.iter().enumerate() {
1191                flat_idx += idx * strides[d];
1192            }
1193            result_data[flat_idx] = src_data[i];
1194        }
1195        Self::from_data(result_data, self_shape.to_vec(), self.device)
1196    }
1197    /// Scatter values to a selected slice along dimension (PyTorch-compatible)
1198    ///
1199    /// Embeds the values of src tensor into self at the given index along dimension dim.
1200    /// This is the inverse of `select()` operation.
1201    ///
1202    /// # PyTorch Compatibility
1203    /// Equivalent to `torch.select_scatter(tensor, src, dim, index)`
1204    ///
1205    /// # Arguments
1206    /// * `src` - Source tensor to scatter (shape should match self with dim removed)
1207    /// * `dim` - Dimension along which to select
1208    /// * `index` - Index position to scatter to
1209    ///
1210    /// # Examples
1211    /// ```ignore
1212    /// let tensor = Tensor::zeros(&[3, 4, 5], DeviceType::Cpu)?;
1213    /// let src = Tensor::ones(&[3, 5], DeviceType::Cpu)?; // dim=1 removed
1214    /// let result = tensor.select_scatter(&src, 1, 2)?;
1215    /// // result[:, 2, :] = src
1216    /// ```
1217    pub fn select_scatter(&self, src: &Tensor<T>, dim: isize, index: isize) -> Result<Self> {
1218        let ndim = self.ndim() as isize;
1219        let dim_normalized = if dim < 0 { ndim + dim } else { dim };
1220        if dim_normalized < 0 || dim_normalized >= ndim {
1221            return Err(TorshError::InvalidArgument(format!(
1222                "Dimension {} out of range for {}-dimensional tensor",
1223                dim,
1224                self.ndim()
1225            )));
1226        }
1227        let dim_u = dim_normalized as usize;
1228        let self_shape_ref = self.shape();
1229        let self_shape = self_shape_ref.dims();
1230        let index_normalized = if index < 0 {
1231            (self_shape[dim_u] as isize) + index
1232        } else {
1233            index
1234        };
1235        if index_normalized < 0 || index_normalized >= self_shape[dim_u] as isize {
1236            return Err(TorshError::InvalidArgument(format!(
1237                "Index {} out of bounds for dimension {} with size {}",
1238                index, dim_u, self_shape[dim_u]
1239            )));
1240        }
1241        let index_u = index_normalized as usize;
1242        let expected_src_shape: Vec<usize> = self_shape
1243            .iter()
1244            .enumerate()
1245            .filter(|(i, _)| *i != dim_u)
1246            .map(|(_, &s)| s)
1247            .collect();
1248        let src_shape_ref = src.shape();
1249        let src_shape = src_shape_ref.dims();
1250        if src_shape != expected_src_shape.as_slice() {
1251            return Err(TorshError::ShapeMismatch {
1252                expected: expected_src_shape,
1253                got: src_shape.to_vec(),
1254            });
1255        }
1256        let mut result_data = self.to_vec()?;
1257        let src_data = src.to_vec()?;
1258        let mut self_strides = vec![1; self_shape.len()];
1259        for i in (0..self_shape.len() - 1).rev() {
1260            self_strides[i] = self_strides[i + 1] * self_shape[i + 1];
1261        }
1262        let outer_size: usize = self_shape[..dim_u].iter().product();
1263        let inner_size: usize = self_shape[dim_u + 1..].iter().product();
1264        for outer in 0..outer_size {
1265            for inner in 0..inner_size {
1266                let self_idx =
1267                    outer * (self_shape[dim_u] * inner_size) + index_u * inner_size + inner;
1268                let src_idx = outer * inner_size + inner;
1269                result_data[self_idx] = src_data[src_idx];
1270            }
1271        }
1272        Self::from_data(result_data, self_shape.to_vec(), self.device)
1273    }
1274    /// Scatter values to a slice along dimension (PyTorch-compatible)
1275    ///
1276    /// Embeds the values of src tensor into self along dimension dim, starting at
1277    /// start index, ending at end index, with the given step.
1278    ///
1279    /// # PyTorch Compatibility
1280    /// Equivalent to `torch.slice_scatter(tensor, src, dim, start, end, step)`
1281    ///
1282    /// # Arguments
1283    /// * `src` - Source tensor to scatter
1284    /// * `dim` - Dimension along which to slice
1285    /// * `start` - Starting index (None means 0)
1286    /// * `end` - Ending index (None means size of dim)
1287    /// * `step` - Step size (default: 1)
1288    ///
1289    /// # Examples
1290    /// ```ignore
1291    /// let tensor = Tensor::zeros(&[5, 5], DeviceType::Cpu)?;
1292    /// let src = Tensor::ones(&[2, 5], DeviceType::Cpu)?;
1293    /// let result = tensor.slice_scatter(&src, 0, Some(1), Some(3), 1)?;
1294    /// // result[1:3, :] = src
1295    /// ```
1296    pub fn slice_scatter(
1297        &self,
1298        src: &Tensor<T>,
1299        dim: isize,
1300        start: Option<isize>,
1301        end: Option<isize>,
1302        step: usize,
1303    ) -> Result<Self> {
1304        if step == 0 {
1305            return Err(TorshError::InvalidArgument(
1306                "Step must be greater than 0".to_string(),
1307            ));
1308        }
1309        let ndim = self.ndim() as isize;
1310        let dim_normalized = if dim < 0 { ndim + dim } else { dim };
1311        if dim_normalized < 0 || dim_normalized >= ndim {
1312            return Err(TorshError::InvalidArgument(format!(
1313                "Dimension {} out of range for {}-dimensional tensor",
1314                dim,
1315                self.ndim()
1316            )));
1317        }
1318        let dim_u = dim_normalized as usize;
1319        let self_shape_ref = self.shape();
1320        let self_shape = self_shape_ref.dims();
1321        let dim_size = self_shape[dim_u] as isize;
1322        let start_normalized = start.unwrap_or(0);
1323        let start_normalized = if start_normalized < 0 {
1324            dim_size + start_normalized
1325        } else {
1326            start_normalized
1327        };
1328        let start_normalized = std::cmp::max(0, std::cmp::min(start_normalized, dim_size)) as usize;
1329        let end_normalized = end.unwrap_or(dim_size);
1330        let end_normalized = if end_normalized < 0 {
1331            dim_size + end_normalized
1332        } else {
1333            end_normalized
1334        };
1335        let end_normalized = std::cmp::max(0, std::cmp::min(end_normalized, dim_size)) as usize;
1336        let slice_len = if end_normalized > start_normalized {
1337            (end_normalized - start_normalized + step - 1) / step
1338        } else {
1339            0
1340        };
1341        let mut expected_src_shape = self_shape.to_vec();
1342        expected_src_shape[dim_u] = slice_len;
1343        let src_shape_ref = src.shape();
1344        let src_shape = src_shape_ref.dims();
1345        if src_shape != expected_src_shape.as_slice() {
1346            return Err(TorshError::ShapeMismatch {
1347                expected: expected_src_shape,
1348                got: src_shape.to_vec(),
1349            });
1350        }
1351        let mut result_data = self.to_vec()?;
1352        let src_data = src.to_vec()?;
1353        let mut self_strides = vec![1; self_shape.len()];
1354        for i in (0..self_shape.len() - 1).rev() {
1355            self_strides[i] = self_strides[i + 1] * self_shape[i + 1];
1356        }
1357        let outer_size: usize = self_shape[..dim_u].iter().product();
1358        let inner_size: usize = self_shape[dim_u + 1..].iter().product();
1359        for outer in 0..outer_size {
1360            for slice_idx in 0..slice_len {
1361                let self_dim_idx = start_normalized + slice_idx * step;
1362                for inner in 0..inner_size {
1363                    let self_idx = outer * (self_shape[dim_u] * inner_size)
1364                        + self_dim_idx * inner_size
1365                        + inner;
1366                    let src_idx = outer * (slice_len * inner_size) + slice_idx * inner_size + inner;
1367                    result_data[self_idx] = src_data[src_idx];
1368                }
1369            }
1370        }
1371        Self::from_data(result_data, self_shape.to_vec(), self.device)
1372    }
1373}
torsh_tensor/data_ops/functions.rs

torsh_tensor/data_ops/
functions.rs
