Fix tensor broadcast off-by-one error.

Caught by JAX unit tests.  Triggered if broadcast is smaller than packet
size.


(cherry picked from commit ffb78e23a1)
This commit is contained in:
Antonio Sanchez
2021-11-15 22:19:04 -08:00
committed by Rasmus Munk Larsen
parent c379a21191
commit 7e3bc4177e

View File

@@ -461,8 +461,8 @@ struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
values[i] = m_impl.coeff(inputIndex);
++outputOffset;
} else {
outputOffset = 0;
values[i] = m_impl.coeff(++inputIndex);
outputOffset = 1; // Next offset.
}
}
return internal::pload<PacketReturnType>(values);