oarust 0.8.6

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=matrix_elemwise.json schema_version=1 generator_version=64
// schema_sha256=6d4112435554d3615e6c2494c5062e5f93fa85815bccffa2ca49b12f97aaac75
//! Elementwise Matrix operations and shared broadcast-shape inference.

use crate::{
	DType, Error, Matrix, OpAttribute, OpShapeRule, Result,
	matrix::{self, autograd, sum_to_shape},
	runtime::{BufferBinding, ComputeDispatch, KernelId, PushConstant, SemanticDispatch},
};

const BROADCAST_RANK_LIMIT: usize = 8;

struct BroadcastLayout {
	output_shape: Vec<usize>,
	output_dims: [u32; BROADCAST_RANK_LIMIT],
	left_strides: [u32; BROADCAST_RANK_LIMIT],
	right_strides: [u32; BROADCAST_RANK_LIMIT],
	element_count: u32,
	rank: u32,
}

fn resolve_broadcast(
	left: &[usize],
	right: &[usize],
	operation: &'static str,
) -> Result<BroadcastLayout> {
	let rank = left.len().max(right.len());
	if rank > BROADCAST_RANK_LIMIT {
		return Err(Error::invalid_argument(format!(
			"{operation} supports at most {BROADCAST_RANK_LIMIT} broadcast axes"
		)));
	}
	let mut output_shape = vec![1_usize; rank];
	for (axis, output_extent) in output_shape.iter_mut().enumerate() {
		let left_extent = aligned_extent(left, rank, axis);
		let right_extent = aligned_extent(right, rank, axis);
		*output_extent = if left_extent == right_extent {
			left_extent
		} else if left_extent == 1 {
			right_extent
		} else if right_extent == 1 {
			left_extent
		} else {
			return Err(Error::invalid_argument(format!(
				"{operation} cannot broadcast shapes {left:?} and {right:?}"
			)));
		};
	}
	let element_count = output_shape.iter().try_fold(1_usize, |product, extent| {
		product
			.checked_mul(*extent)
			.ok_or_else(|| Error::invalid_argument(format!("{operation} output size overflows usize")))
	})?;
	let element_count = u32::try_from(element_count)
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let mut output_dims = [1_u32; BROADCAST_RANK_LIMIT];
	for (destination, extent) in output_dims.iter_mut().zip(&output_shape) {
		*destination = u32::try_from(*extent)
			.map_err(|_| Error::invalid_argument(format!("{operation} output extent exceeds u32")))?;
	}
	Ok(BroadcastLayout {
		left_strides: broadcast_strides(left, &output_shape, operation)?,
		right_strides: broadcast_strides(right, &output_shape, operation)?,
		output_shape,
		output_dims,
		element_count,
		rank: u32::try_from(rank).expect("broadcast rank limit fits u32"),
	})
}

fn broadcast_strides(
	shape: &[usize],
	output_shape: &[usize],
	operation: &'static str,
) -> Result<[u32; BROADCAST_RANK_LIMIT]> {
	let rank = output_shape.len();
	let offset = rank - shape.len();
	let mut strides = [0_u32; BROADCAST_RANK_LIMIT];
	let mut contiguous_stride = 1_usize;
	for source_axis in (0..shape.len()).rev() {
		let axis = offset + source_axis;
		if shape[source_axis] != 1 || output_shape[axis] == 1 {
			strides[axis] = u32::try_from(contiguous_stride)
				.map_err(|_| Error::invalid_argument(format!("{operation} input stride exceeds u32")))?;
		}
		contiguous_stride = contiguous_stride
			.checked_mul(shape[source_axis])
			.ok_or_else(|| Error::invalid_argument(format!("{operation} input size overflows")))?;
	}
	Ok(strides)
}

fn aligned_extent(shape: &[usize], rank: usize, axis: usize) -> usize {
	axis
		.checked_sub(rank - shape.len())
		.map_or(1, |source_axis| shape[source_axis])
}

fn select_kernel(
	dtype: DType,
	routes: &[(DType, KernelId)],
	operation: &'static str,
) -> Result<KernelId> {
	routes
		.iter()
		.find_map(|(candidate, kernel)| (*candidate == dtype).then_some(*kernel))
		.ok_or_else(|| {
			Error::invalid_argument(format!(
				"{operation} does not support dtype {}",
				dtype.token()
			))
		})
}

/// Add two matrices elementwise using multidirectional broadcasting.
///
/// Shapes use multidirectional broadcasting over at most eight axes.
///
/// The I32 route uses two's-complement wrapping arithmetic.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn add(left: &Matrix, right: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::ADD;
	let operation = contract.name();
	if left.dtype() != right.dtype() {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal dtypes; left is {}, right is {}",
			left.dtype().token(),
			right.dtype().token()
		)));
	}
	let engine = left.engine_handle();
	if !engine.same_as(right.engine_handle()) {
		return Err(Error::invalid_argument(format!(
			"{operation} inputs must belong to the same engine"
		)));
	}
	let output = if left.shape() == right.shape() {
		let kernel = select_kernel(
			left.dtype(),
			&[
				(DType::F32, KernelId::MatrixAddF32),
				(DType::I32, KernelId::MatrixAddI32),
			],
			operation,
		)?;
		let element_count = u32::try_from(left.element_count())
			.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
		let output = Matrix::allocate(
			engine,
			left.shape().to_vec(),
			left.element_count(),
			left.dtype(),
		)?;
		if element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants = [PushConstant::U32(element_count)];
			let inputs = [left, right];
			let outputs = [&output];
			engine.record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else if contract.shape_rule() == OpShapeRule::Broadcast {
		let operation = contract.name();
		let layout = resolve_broadcast(left.shape(), right.shape(), operation)?;
		let kernel = select_kernel(
			left.dtype(),
			&[
				(DType::F32, KernelId::MatrixAddBroadcastF32),
				(DType::I32, KernelId::MatrixAddBroadcastI32),
			],
			operation,
		)?;
		let output_count = usize::try_from(layout.element_count)
			.expect("u32 element count always fits the supported Rust targets");
		let output = Matrix::allocate(
			left.engine_handle(),
			layout.output_shape,
			output_count,
			left.dtype(),
		)?;
		if layout.element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants: [PushConstant; 26] = std::array::from_fn(|index| {
				PushConstant::U32(match index {
					0 => layout.element_count,
					1 => layout.rank,
					2..=9 => layout.output_dims[index - 2],
					10..=17 => layout.left_strides[index - 10],
					_ => layout.right_strides[index - 18],
				})
			});
			let inputs = [left, right];
			let outputs = [&output];
			left.engine_handle().record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(layout.element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal shapes; left is {:?}, right is {:?}",
			left.shape(),
			right.shape()
		)));
	};
	if left.dtype() == DType::F32 {
		autograd::record_add(left, right, &output)?;
	}
	Ok(output)
}

/// Subtract two FP32 matrices elementwise using multidirectional broadcasting.
///
/// Shapes use multidirectional broadcasting over at most eight axes.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn sub(left: &Matrix, right: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::SUB;
	let operation = contract.name();
	if left.dtype() != right.dtype() {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal dtypes; left is {}, right is {}",
			left.dtype().token(),
			right.dtype().token()
		)));
	}
	let engine = left.engine_handle();
	if !engine.same_as(right.engine_handle()) {
		return Err(Error::invalid_argument(format!(
			"{operation} inputs must belong to the same engine"
		)));
	}
	let output = if left.shape() == right.shape() {
		let kernel = select_kernel(
			left.dtype(),
			&[(DType::F32, KernelId::MatrixSubF32)],
			operation,
		)?;
		let element_count = u32::try_from(left.element_count())
			.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
		let output = Matrix::allocate(
			engine,
			left.shape().to_vec(),
			left.element_count(),
			left.dtype(),
		)?;
		if element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants = [PushConstant::U32(element_count)];
			let inputs = [left, right];
			let outputs = [&output];
			engine.record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else if contract.shape_rule() == OpShapeRule::Broadcast {
		let operation = contract.name();
		let layout = resolve_broadcast(left.shape(), right.shape(), operation)?;
		let kernel = select_kernel(
			left.dtype(),
			&[(DType::F32, KernelId::MatrixSubBroadcastF32)],
			operation,
		)?;
		let output_count = usize::try_from(layout.element_count)
			.expect("u32 element count always fits the supported Rust targets");
		let output = Matrix::allocate(
			left.engine_handle(),
			layout.output_shape,
			output_count,
			left.dtype(),
		)?;
		if layout.element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants: [PushConstant; 26] = std::array::from_fn(|index| {
				PushConstant::U32(match index {
					0 => layout.element_count,
					1 => layout.rank,
					2..=9 => layout.output_dims[index - 2],
					10..=17 => layout.left_strides[index - 10],
					_ => layout.right_strides[index - 18],
				})
			});
			let inputs = [left, right];
			let outputs = [&output];
			left.engine_handle().record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(layout.element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal shapes; left is {:?}, right is {:?}",
			left.shape(),
			right.shape()
		)));
	};
	if left.dtype() == DType::F32 {
		autograd::record_sub(left, right, &output)?;
	}
	Ok(output)
}

/// Multiply two FP32 matrices elementwise.
///
/// Shapes use multidirectional broadcasting over at most eight axes.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn mul(left: &Matrix, right: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::MUL;
	let operation = contract.name();
	if left.dtype() != right.dtype() {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal dtypes; left is {}, right is {}",
			left.dtype().token(),
			right.dtype().token()
		)));
	}
	let engine = left.engine_handle();
	if !engine.same_as(right.engine_handle()) {
		return Err(Error::invalid_argument(format!(
			"{operation} inputs must belong to the same engine"
		)));
	}
	let output = if left.shape() == right.shape() {
		let kernel = select_kernel(
			left.dtype(),
			&[(DType::F32, KernelId::MatrixMulF32)],
			operation,
		)?;
		let element_count = u32::try_from(left.element_count())
			.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
		let output = Matrix::allocate(
			engine,
			left.shape().to_vec(),
			left.element_count(),
			left.dtype(),
		)?;
		if element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants = [PushConstant::U32(element_count)];
			let inputs = [left, right];
			let outputs = [&output];
			engine.record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else if contract.shape_rule() == OpShapeRule::Broadcast {
		let operation = contract.name();
		let layout = resolve_broadcast(left.shape(), right.shape(), operation)?;
		let kernel = select_kernel(
			left.dtype(),
			&[(DType::F32, KernelId::MatrixMulBroadcastF32)],
			operation,
		)?;
		let output_count = usize::try_from(layout.element_count)
			.expect("u32 element count always fits the supported Rust targets");
		let output = Matrix::allocate(
			left.engine_handle(),
			layout.output_shape,
			output_count,
			left.dtype(),
		)?;
		if layout.element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants: [PushConstant; 26] = std::array::from_fn(|index| {
				PushConstant::U32(match index {
					0 => layout.element_count,
					1 => layout.rank,
					2..=9 => layout.output_dims[index - 2],
					10..=17 => layout.left_strides[index - 10],
					_ => layout.right_strides[index - 18],
				})
			});
			let inputs = [left, right];
			let outputs = [&output];
			left.engine_handle().record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(layout.element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal shapes; left is {:?}, right is {:?}",
			left.shape(),
			right.shape()
		)));
	};
	if left.dtype() == DType::F32 {
		autograd::record_mul(left, right, &output)?;
	}
	Ok(output)
}

/// Divide two FP32 matrices elementwise using multidirectional broadcasting.
///
/// Shapes use multidirectional broadcasting over at most eight axes.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn div(left: &Matrix, right: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::DIV;
	let operation = contract.name();
	if left.dtype() != right.dtype() {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal dtypes; left is {}, right is {}",
			left.dtype().token(),
			right.dtype().token()
		)));
	}
	let engine = left.engine_handle();
	if !engine.same_as(right.engine_handle()) {
		return Err(Error::invalid_argument(format!(
			"{operation} inputs must belong to the same engine"
		)));
	}
	let output = if left.shape() == right.shape() {
		let kernel = select_kernel(
			left.dtype(),
			&[(DType::F32, KernelId::MatrixDivF32)],
			operation,
		)?;
		let element_count = u32::try_from(left.element_count())
			.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
		let output = Matrix::allocate(
			engine,
			left.shape().to_vec(),
			left.element_count(),
			left.dtype(),
		)?;
		if element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants = [PushConstant::U32(element_count)];
			let inputs = [left, right];
			let outputs = [&output];
			engine.record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else if contract.shape_rule() == OpShapeRule::Broadcast {
		let operation = contract.name();
		let layout = resolve_broadcast(left.shape(), right.shape(), operation)?;
		let kernel = select_kernel(
			left.dtype(),
			&[(DType::F32, KernelId::MatrixDivBroadcastF32)],
			operation,
		)?;
		let output_count = usize::try_from(layout.element_count)
			.expect("u32 element count always fits the supported Rust targets");
		let output = Matrix::allocate(
			left.engine_handle(),
			layout.output_shape,
			output_count,
			left.dtype(),
		)?;
		if layout.element_count != 0 {
			let buffers = [
				BufferBinding::read(left.storage()),
				BufferBinding::read(right.storage()),
				BufferBinding::write(output.storage()),
			];
			let push_constants: [PushConstant; 26] = std::array::from_fn(|index| {
				PushConstant::U32(match index {
					0 => layout.element_count,
					1 => layout.rank,
					2..=9 => layout.output_dims[index - 2],
					10..=17 => layout.left_strides[index - 10],
					_ => layout.right_strides[index - 18],
				})
			});
			let inputs = [left, right];
			let outputs = [&output];
			left.engine_handle().record_semantic(
				ComputeDispatch {
					kernel,
					buffers: &buffers,
					push_constants: &push_constants,
					workgroups: kernel.linear_workgroups(layout.element_count),
				},
				SemanticDispatch {
					contract,
					inputs: &inputs,
					outputs: &outputs,
					attributes: &[],
				},
			)?;
		}
		output
	} else {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal shapes; left is {:?}, right is {:?}",
			left.shape(),
			right.shape()
		)));
	};
	if left.dtype() == DType::F32 {
		autograd::record_div(left, right, &output)?;
	}
	Ok(output)
}

/// Multiply every FP32 matrix element by a scalar.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn scale(input: &Matrix, scalar: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::SCALE;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixScaleF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count), PushConstant::F32(scalar)];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(scalar),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}
	autograd::record_scale(input, &output, scalar)?;
	Ok(output)
}

/// Negate every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn neg(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::NEG;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixNegF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_scale(input, &output, -1.0)?;
	Ok(output)
}

/// Take the absolute value of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn abs(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::ABS;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixAbsF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_abs(input, &output)?;
	Ok(output)
}

/// Take the natural logarithm of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn log(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::LOG;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixLogF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_log(input, &output)?;
	Ok(output)
}

/// Take the square root of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn sqrt(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::SQRT;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixSqrtF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_sqrt(input, &output)?;
	Ok(output)
}

/// Raise every FP32 matrix element to a scalar exponent.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn pow(input: &Matrix, exponent: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::POW;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixPowF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [
			PushConstant::U32(element_count),
			PushConstant::F32(exponent),
		];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(exponent),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}

	Ok(output)
}

/// Add a scalar to every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn add_scalar(input: &Matrix, scalar: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::ADD_SCALAR;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixAddScalarF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count), PushConstant::F32(scalar)];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(scalar),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}
	autograd::record_scale(input, &output, 1.0)?;
	Ok(output)
}

/// Subtract a scalar from every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn sub_scalar(input: &Matrix, scalar: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::SUB_SCALAR;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixSubScalarF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count), PushConstant::F32(scalar)];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(scalar),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}
	autograd::record_scale(input, &output, 1.0)?;
	Ok(output)
}

/// Divide every FP32 matrix element by a scalar.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn div_scalar(input: &Matrix, scalar: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::DIV_SCALAR;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixDivScalarF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count), PushConstant::F32(scalar)];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(scalar),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}
	autograd::record_scale(input, &output, 1.0 / scalar)?;
	Ok(output)
}

/// Take the exponential of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn exp(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::EXP;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixExpF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_exp(input, &output)?;
	Ok(output)
}

/// Take the sine of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn sin(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::SIN;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixSinF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}

	Ok(output)
}

/// Take the cosine of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn cos(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::COS;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixCosF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}

	Ok(output)
}

/// Take the reciprocal of every FP32 matrix element.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn reciprocal(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::RECIPROCAL;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixReciprocalF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_reciprocal(input, &output)?;
	Ok(output)
}

/// Cap every FP32 matrix element at an upper bound.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn clamp_max(input: &Matrix, maximum: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::CLAMP_MAX;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixClampMaxF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count), PushConstant::F32(maximum)];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(maximum),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}
	autograd::record_clamp_max(input, &output, maximum)?;
	Ok(output)
}

/// Raise every FP32 matrix element to a lower bound.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn clamp_min(input: &Matrix, minimum: f32) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::CLAMP_MIN;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[(DType::F32, KernelId::MatrixClampMinF32)],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count), PushConstant::F32(minimum)];
		let inputs = [input];
		let outputs = [&output];
		let attributes = [OpAttribute::Float {
			name: contract.attribute_specs()[0].name().into(),
			value: f64::from(minimum),
		}];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &attributes,
			},
		)?;
	}
	autograd::record_clamp_min(input, &output, minimum)?;
	Ok(output)
}

/// Deep-copy a dense matrix into independent device storage.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.
pub fn copy(input: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::COPY;
	let operation = contract.name();
	let kernel = select_kernel(
		input.dtype(),
		&[
			(DType::F32, KernelId::MatrixCopyF32),
			(DType::U8, KernelId::MatrixCopyU8),
			(DType::I32, KernelId::MatrixCopyI32),
			(DType::U32, KernelId::MatrixCopyU32),
		],
		operation,
	)?;
	let engine = input.engine_handle();
	let element_count = u32::try_from(input.element_count())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.element_count(),
		input.dtype(),
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(input.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		let inputs = [input];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
	}
	autograd::record_copy(input, &output)?;
	Ok(output)
}

/// Compute the copy adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn copy_backward(output_gradient: Matrix) -> Result<Matrix> {
	Ok(output_gradient)
}

/// Compute the add adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn add_backward(
	left: &Matrix,
	right: &Matrix,
	output_gradient: &Matrix,
) -> Result<(Matrix, Matrix)> {
	let left_gradient = sum_to_shape(output_gradient, left.shape())?;
	let right_gradient = sum_to_shape(output_gradient, right.shape())?;
	Ok((left_gradient, right_gradient))
}

/// Compute the mul adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn mul_backward(
	left: &Matrix,
	right: &Matrix,
	output_gradient: &Matrix,
) -> Result<(Matrix, Matrix)> {
	let left_gradient = matrix::mul(output_gradient, right)?;
	let right_gradient = matrix::mul(output_gradient, left)?;
	let left_gradient = sum_to_shape(&left_gradient, left.shape())?;
	let right_gradient = sum_to_shape(&right_gradient, right.shape())?;
	Ok((left_gradient, right_gradient))
}

/// Compute the div adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn div_backward(
	left: &Matrix,
	right: &Matrix,
	output_gradient: &Matrix,
) -> Result<(Matrix, Matrix)> {
	let left_gradient = matrix::div(output_gradient, right)?;
	let denominator = matrix::mul(right, right)?;
	let right_gradient = matrix::scale(
		&matrix::div(&matrix::mul(output_gradient, left)?, &denominator)?,
		-1.0,
	)?;
	let left_gradient = sum_to_shape(&left_gradient, left.shape())?;
	let right_gradient = sum_to_shape(&right_gradient, right.shape())?;
	Ok((left_gradient, right_gradient))
}

/// Compute the scale adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn scale_backward(output_gradient: &Matrix, scalar: f32) -> Result<Matrix> {
	matrix::scale(output_gradient, scalar)
}

/// Compute the reciprocal adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn reciprocal_backward(output: &Matrix, output_gradient: &Matrix) -> Result<Matrix> {
	let output_squared = matrix::mul(output, output)?;
	let input_gradient = matrix::scale(&matrix::mul(output_gradient, &output_squared)?, -1.0)?;
	Ok(input_gradient)
}

/// Compute the exp adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn exp_backward(output: &Matrix, output_gradient: &Matrix) -> Result<Matrix> {
	let input_gradient = matrix::mul(output_gradient, output)?;
	Ok(input_gradient)
}

/// Compute the log adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn log_backward(input: &Matrix, output_gradient: &Matrix) -> Result<Matrix> {
	let input_gradient = matrix::div(output_gradient, input)?;
	Ok(input_gradient)
}

/// Compute the abs adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn abs_backward(input: &Matrix, output_gradient: &Matrix) -> Result<Matrix> {
	let positive = matrix::equal(&matrix::clamp_max(input, 0.0)?, 0.0)?;
	let negative = matrix::equal(&matrix::clamp_min(input, 0.0)?, 0.0)?;
	let sign = matrix::sub(&positive, &negative)?;
	let input_gradient = matrix::mul(output_gradient, &sign)?;
	Ok(input_gradient)
}

/// Compute the sqrt adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn sqrt_backward(output: &Matrix, output_gradient: &Matrix) -> Result<Matrix> {
	let input_gradient = matrix::div(output_gradient, &matrix::scale(output, 2.0)?)?;
	Ok(input_gradient)
}

/// Compute the clamp_max adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn clamp_max_backward(
	input: &Matrix,
	output_gradient: &Matrix,
	maximum: f32,
) -> Result<Matrix> {
	let difference = matrix::sub_scalar(input, maximum)?;
	let mask = matrix::equal(&matrix::clamp_min(&difference, 0.0)?, 0.0)?;
	let input_gradient = matrix::mul(output_gradient, &mask)?;
	Ok(input_gradient)
}

/// Compute the clamp_min adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn clamp_min_backward(
	input: &Matrix,
	output_gradient: &Matrix,
	minimum: f32,
) -> Result<Matrix> {
	let difference = matrix::sub_scalar(input, minimum)?;
	let mask = matrix::equal(&matrix::clamp_max(&difference, 0.0)?, 0.0)?;
	let input_gradient = matrix::mul(output_gradient, &mask)?;
	Ok(input_gradient)
}

/// Compute the sub adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn sub_backward(
	left: &Matrix,
	right: &Matrix,
	output_gradient: &Matrix,
) -> Result<(Matrix, Matrix)> {
	let left_gradient = sum_to_shape(output_gradient, left.shape())?;
	let right_gradient = sum_to_shape(&matrix::scale(output_gradient, -1.0)?, right.shape())?;
	Ok((left_gradient, right_gradient))
}