Skip to content
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

Range on TDIM output i64 instead of TDIM #1604

Merged
merged 2 commits into from
Dec 12, 2024
Merged

Range on TDIM output i64 instead of TDIM #1604

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b1c73c4
Range on TDIM output i64 instead of TDIM
hubertdelajonquieresonos Dec 11, 2024
bb1e4c5
Fix Range declutter
hubertdelajonquieresonos Dec 12, 2024
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
41 changes: 25 additions & 16 deletions core/src/ops/array/range.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::ops::cast::Cast;
use tract_num_traits::AsPrimitive;
use tract_num_traits::Zero;

Expand Down Expand Up @@ -59,14 +60,14 @@ impl Range {
values: &SymbolValues,
) -> TractResult<Tensor> {
if start.datum_type() == TDim::datum_type() {
let start = start.to_scalar::<TDim>()?.eval(values).to_i64()?;
let step = step.to_scalar::<TDim>()?.eval(values).to_i64()?;
let len = {
let start = start.to_scalar::<TDim>()?.eval(values).to_i64()?;
let end = end.to_scalar::<TDim>()?.eval(values).to_i64()?;
let step = step.to_scalar::<TDim>()?.eval(values).to_i64()?;
#[allow(clippy::cast_abs_to_unsigned)]
((end - start).abs() as usize).divceil(step.abs() as usize)
};
Self::make_t::<TDim>(start, step, len)
Self::make_t::<i64>(&tensor0(start), &tensor0(step), len)
} else {
let len = dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
self, start, end, step
Expand All @@ -90,18 +91,25 @@ impl Range {

impl TypedOp for Range {
fn declutter(
&self,
model: &TypedModel,
node: &TypedNode,
) -> TractResult<Option<TypedModelPatch>> {
&self,
model: &TypedModel,
node: &TypedNode,
) -> TractResult<Option<TypedModelPatch>> {
let Some(succ) = model.single_succ(node.id)? else { return Ok(None) };
let Some(slice) = succ.op_as::<Slice>() else { return Ok(None) };
if slice.start.is_zero() && slice.end.is_one() {
let mut patch = TypedModelPatch::default();
let wire = patch.tap_model(model, node.inputs[0])?;
let mut wire = patch.tap_model(model, node.inputs[0])?;
if model.outlet_fact(node.inputs[0])?.datum_type.is_tdim() {
wire = patch.wire_node(
format!("{}.cast-tdim", node.name),
Cast { to: DatumType::I64 },
&[wire],
)?[0];
}
let wire = patch.wire_node(&node.name, AxisOp::Add(0), &[wire])?;
patch.shunt_outside(model, succ.id.into(), wire[0])?;
return Ok(Some(patch))
return Ok(Some(patch));
}
Ok(None)
}
Expand All @@ -116,22 +124,23 @@ impl TypedOp for Range {
ensure!(end.shape.volume().is_one());
ensure!(step.shape.volume().is_one());
if let (Some(start), Some(end), Some(step)) = (&start.konst, &end.konst, &step.konst) {
let len = if start.datum_type() == TDim::datum_type() {
if start.datum_type() == TDim::datum_type() {
let start = start.to_scalar::<TDim>()?;
let end = end.to_scalar::<TDim>()?;
let step = step.cast_to_scalar::<i64>()?;
if step < 0 {
let len = if step < 0 {
(start.clone() - end).divceil(-step as usize)
} else {
(end.clone() - start).divceil(step as usize)
}
};
Ok(tvec!(DatumType::I64.fact([len])))
} else {
dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
let len = dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
self, start, end, step
))?
.to_dim()
};
Ok(tvec!(start.datum_type().fact([len])))
.to_dim();
Ok(tvec!(start.datum_type().fact([len])))
}
} else {
Ok(tvec!(start.datum_type.fact(&[self.len.clone()])))
}
Expand Down
Loading