Skip to content

Fix: broadcast during slice assign and align assign API #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CodeWithKyrian merged 1 commit into from
May 13, 2026
Merged

Fix: broadcast during slice assign and align assign API #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion rust/src/ffi/indexing/assign.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,15 @@
//!
//! Provides assign operations between strided array views.

use crate::helpers::error::{set_last_error, ERR_GENERIC, SUCCESS};
use crate::helpers::error::{set_last_error, ERR_GENERIC, ERR_SHAPE, SUCCESS};
use crate::helpers::view::{
extract_view_bool, extract_view_c128, extract_view_c64, extract_view_f32, extract_view_f64,
extract_view_i16, extract_view_i32, extract_view_i64, extract_view_i8, extract_view_mut_bool,
extract_view_mut_c128, extract_view_mut_c64, extract_view_mut_f32, extract_view_mut_f64,
extract_view_mut_i16, extract_view_mut_i32, extract_view_mut_i64, extract_view_mut_i8,
extract_view_mut_u16, extract_view_mut_u32, extract_view_mut_u64, extract_view_mut_u8,
extract_view_u16, extract_view_u32, extract_view_u64, extract_view_u8,
rhs_broadcasts_to_lhs,
};
use crate::types::dtype::DType;
use crate::types::{ArrayMetadata, NdArrayHandle};
Expand Down Expand Up @@ -48,6 +49,16 @@ pub unsafe extern "C" fn ndarray_assign(
return ERR_GENERIC;
}

let dst_shape = unsafe { dst_meta.shape_slice() };
let src_shape = unsafe { src_meta.shape_slice() };
if !rhs_broadcasts_to_lhs(dst_shape, src_shape) {
set_last_error(format!(
"Cannot assign: source shape {:?} cannot broadcast to destination shape {:?}",
src_shape, dst_shape
));
return ERR_SHAPE;
}

let is_same = dst_wrapper.is_same_array(src_wrapper);

match dst_wrapper.dtype {
Expand Down
40 changes: 40 additions & 0 deletions rust/src/helpers/view.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -404,3 +404,43 @@ pub fn broadcast_shape(shape_a: &[usize], shape_b: &[usize]) -> Option<Vec<usize
}
Some(out)
}

/// Whether `src` can be broadcast to `dst` for in-place assignment (`rhs` broadcast to `lhs` shape).
///
/// Matches ndarray `ArrayBase::assign`: the right-hand side is broadcast to the destination shape;
/// this returns `false` when that broadcast is not defined.
pub fn rhs_broadcasts_to_lhs(dst: &[usize], src: &[usize]) -> bool {
match broadcast_shape(dst, src) {
Some(bc) => bc == dst,
None => false,
}
}

#[cfg(test)]
mod broadcast_tests {
use super::{rhs_broadcasts_to_lhs, broadcast_shape};

#[test]
fn rhs_broadcasts_to_lhs_matches_elementwise_broadcast() {
assert!(rhs_broadcasts_to_lhs(&[3, 4], &[3, 1]));
assert!(rhs_broadcasts_to_lhs(&[3, 4], &[1, 4]));
assert!(rhs_broadcasts_to_lhs(&[3], &[1]));
assert!(rhs_broadcasts_to_lhs(&[1, 3], &[3]));
assert!(rhs_broadcasts_to_lhs(&[3, 4], &[]));
assert!(rhs_broadcasts_to_lhs(&[], &[]));
}

#[test]
fn rhs_broadcasts_to_lhs_rejects_incompatible() {
assert!(!rhs_broadcasts_to_lhs(&[3], &[2]));
assert!(!rhs_broadcasts_to_lhs(&[3], &[1, 3]));
assert!(!rhs_broadcasts_to_lhs(&[3, 4], &[2, 4]));
assert!(!rhs_broadcasts_to_lhs(&[2, 1, 3], &[2, 3]));
}

#[test]
fn broadcast_shape_symmetric_examples() {
assert_eq!(broadcast_shape(&[3, 4], &[3, 1]), Some(vec![3, 4]));
assert_eq!(broadcast_shape(&[3], &[2]), None);
}
}
6 changes: 3 additions & 3 deletions src/Functions.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1639,11 +1639,11 @@ function slice(NDArray $a, array $selection): NDArray
/**
* Assign values to the current array/view.
*
* Supports scalar assignment (fill) or array assignment (from PHP array or NDArray).
* Supports scalar assignment (fill) or NDArray assignment (rhs is broadcast to this view’s shape when compatible).
*
* @param mixed $value Scalar value or array/NDArray
* @param bool|Complex|float|int|NDArray $value Scalar value or NDArray
*/
function assign(NDArray $a, mixed $value): void
function assign(NDArray $a, bool|Complex|float|int|NDArray $value): void
{
$a->assign($value);
}
Expand Down
76 changes: 13 additions & 63 deletions src/Traits/HasSlicing.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
namespace PhpMlKit\NDArray\Traits;

use PhpMlKit\NDArray\ArrayMetadata;
use PhpMlKit\NDArray\Complex;
use PhpMlKit\NDArray\Exceptions\IndexException;
use PhpMlKit\NDArray\Exceptions\ShapeException;
use PhpMlKit\NDArray\FFI\Lib;
use PhpMlKit\NDArray\NDArray;
use PhpMlKit\NDArray\Slice;
Expand Down Expand Up @@ -121,76 +121,26 @@ public function slice(array $selection): self
/**
* Assign values to the current array/view.
*
* Supports scalar assignment (fill) or array assignment (from PHP array or NDArray).
* Supports scalar assignment (fill) or NDArray assignment (rhs is broadcast to this view’s shape when compatible).
*
* @param mixed $value Scalar value or array/NDArray
* @param bool|Complex|float|int|NDArray $value Scalar value or NDArray
*/
public function assign(mixed $value): void
{
if (\is_scalar($value)) {
$this->fill($value);

return;
}

if ($value instanceof NDArray) {
$this->assignFromNDArray($value);

return;
}

throw new \InvalidArgumentException(
'Assignment value must be scalar or NDArray, got '.get_debug_type($value)
);
}

/**
* Fill the array with a scalar value.
*
* @param mixed $value Scalar value
*/
private function fill(mixed $value): void
public function assign(bool|Complex|float|int|NDArray $value): void
{
$ffi = Lib::get();
$cValue = $this->dtype->createCValue($value);

$meta = $this->meta()->toCData();
$status = $ffi->ndarray_fill($this->handle, Lib::addr($meta), Lib::addr($cValue));
$src = $value;

Lib::checkStatus($status);
}
if ($value instanceof NDArray) {
$srcMeta = $src->meta()->toCData();
$dstMeta = $this->meta()->toCData();

/**
* Assign values from an NDArray to the current array/view.
*
* @param NDArray $value Source NDArray
*/
private function assignFromNDArray(NDArray $value): void
{
if ($value->size() !== $this->size()) {
throw new ShapeException(
"Cannot assign array of size {$value->size()} to view of size {$this->size()}"
);
}
$status = $ffi->ndarray_assign($this->handle, Lib::addr($dstMeta), $src->handle(), Lib::addr($srcMeta));
} else {
$cValue = $this->dtype->createCValue($src);

$ffi = Lib::get();
$dstMeta = $this->meta()->toCData();
$src = $value;
if ($src->shape() !== $this->shape()) {
if ($src->isContiguous()) {
$src = $src->reshape($this->shape());
} else {
$src = $src->copy()->reshape($this->shape());
}
$meta = $this->meta()->toCData();
$status = $ffi->ndarray_fill($this->handle, Lib::addr($meta), Lib::addr($cValue));
}
$srcMeta = $src->meta()->toCData();

$status = $ffi->ndarray_assign(
$this->handle,
Lib::addr($dstMeta),
$src->handle(),
Lib::addr($srcMeta)
);

Lib::checkStatus($status);
}
Expand Down
37 changes: 25 additions & 12 deletions tests/Unit/SlicingTest.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,9 +389,32 @@ public function testAssignShapeMismatchThrows(): void
$arr = NDArray::zeros([5]);

$this->expectException(ShapeException::class);
$this->expectExceptionMessage('Cannot assign array of size 2 to view of size 3');
$this->expectExceptionMessage('Cannot assign: source shape');

$arr->slice(['0:3'])->assign(NDArray::array([1, 2]));
$arr->slice(['0:3'])->assign(NDArray::array([1.0, 2.0]));
}

public function testAssignBroadcastsSmallerSourceToSlice(): void
{
$arr = NDArray::zeros([3, 4]);
$arr->slice([':'])->assign(NDArray::ones([1, 4]));

$this->assertSame([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], $arr->toArray());
}

public function testAssignBroadcastsToColumnSlice(): void
{
$arr = NDArray::zeros([3, 4]);
$arr->slice([':', '1:2'])->assign(NDArray::ones([1, 1]));

$this->assertSame(
[
[0.0, 1.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
],
$arr->toArray()
);
}

public function testSliceStepZeroThrows(): void
Expand All @@ -414,16 +437,6 @@ public function testInvalidSelectorThrows(): void
$slice = $arr['invalid:selector'];
}

public function testAssignPHPArrayThrows(): void
{
$arr = NDArray::zeros([5]);

$this->expectException(\InvalidArgumentException::class);
$this->expectExceptionMessage('Assignment value must be scalar or NDArray');

$arr->slice(['0:3'])->assign([1, 2, 3]);
}

public function testSlice3DBasic(): void
{
$arr = NDArray::arange(8)->reshape([2, 2, 2]);
Expand Down
Loading